cve-scripts.py

import argparse
import os
from dotenv import load_dotenv
import csv
import sys

from database import CVEDatabase, HashDatabase

def search_file(db: CVEDatabase | HashDatabase, file_path: str) -> list[tuple[str, str, list]]:
    print(f'Searching file {file_path}')
    matches = []
    with open(file_path, 'r') as file:
        for line in file:
            package, subject = line.strip().split()
            print(f'Querying line {package} {subject}...')
            matches.append(db.query_database(package, subject))

    return matches

def main():
    parser = argparse.ArgumentParser(description='Program designed to query software information for potential vulnerabilities')
    parser.add_argument('--populate', action='store_true', help='Populate the database')
    parser.add_argument('--update', action='store_true', help='Update the database')
    parser.add_argument('--clear', action='store_true', help='Clear the database')
    parser.add_argument('--search', type=str, help='Search the database', nargs=2, metavar=('PACKAGE', 'VERSION/HASH'))
    parser.add_argument('--search-file', type=str, help='Search the database using a file')
    parser.add_argument('--mongo-url', type=str, help='MongoDB URL', default='mongodb://localhost:27017')
    parser.add_argument('--hash-url', type=str, help='Hash Database URL', default='https://hashlookup.circl.lu/')
    parser.add_argument('--nist-api-key', type=str, help='NIST API Key')
    parser.add_argument('--output-file', type=str, help='Report output file', default='report.txt')
    parser.add_argument('--hash-mode', action='store_true', help='Enable hash mode')
    args = parser.parse_args()

    print('CVE Scripts')

    if len(sys.argv) == 1:
        parser.print_help()
        return

    load_dotenv()

    nist_api_key_env = os.getenv('NIST_API_KEY')
    mongo_db_url_env = os.getenv('MONGO_DB_URL')
    hash_db_url_env = os.getenv('HASH_DB_URL')

    db = CVEDatabase(mongo_db_url_env if mongo_db_url_env else args.mongo_url, nist_api_key_env if nist_api_key_env else args.nist_api_key) \
        if not args.hash_mode else HashDatabase(hash_db_url_env if hash_db_url_env else args.hash_url, None)

    print('Connecting to database...')

    if not db.connect():
        print('Is a MongoDB instance running? if yes, check the URL')
        return

    if args.populate:
        db.populate()
    elif args.update:
        db.update()
    elif args.clear:
        db.clear_database()

    if not db.is_populated():
        print('Database is not populated. Please populate the database using the --populate flag.')
        return

    if not db.is_up_to_date():
        print('Database is outdated by more than a week. Please consider updating the database using the --update flag.')

    matches = []

    if args.search:
        package, version = args.search
        print(f'Querying search {package} {version}...')
        matches.append(db.query_database(package, version))

    if args.search_file:
        matches.extend(search_file(db, args.search_file))

    with open(args.output_file, 'w') as file:
        csv_writer = csv.writer(file)
        csv_writer.writerow(['Package', 'Version' if not args.hash_mode else 'Hash', 'Matches'])
        csv_writer.writerows(matches)
        print(f'Report {args.output_file} generated successfully.')

    db.close()

    print('Exiting...')

if __name__ == '__main__':
    main()