database.py

import requests
import time
import pymongo
from pymongo import MongoClient
import datetime
from dateutil import parser
from pyhashlookup.api import Hashlookup

class Database:
    def __init__(self, database_url: str):
        self.database_url: str = database_url

    def populate(self):
        pass

    def update(self):
        pass

    def clear_database(self):
        pass

    def is_populated(self) -> bool:
        return True

    def is_up_to_date(self) -> bool:
        return True

    def close(self):
        pass

class CVEDatabase(Database):
    def __init__(self, mongodb_url: str, api_key: str | None):
        super().__init__(mongodb_url)
        self.api_key = api_key
        self.client = MongoClient(self.database_url, connectTimeoutMS=10000, socketTimeoutMS=10000)
        self.cve_col = self.client['cve-database']['cves']
        self.info_col = self.client['cve-database']['info']

    def connect(self) -> bool:
        try:
            self.info_col.find_one() # Test connection
            return True
        except Exception as e:
            print(f"Error connecting to MongoDB: {e}")
            return False

    def query_database(self, package: str,version: str) -> tuple[str, str, list[dict]]:
        query = {
            'cve.configurations.nodes.cpeMatch.criteria': {
                '$regex': f'.*{package}:{package}.*',
                '$options': 'i'
            }
        }

        return (
            package,
            version,
            [doc for doc in self.cve_col.find(query)]
        )

    def close(self):
        self.client.close()

    def populate(self):
        print('Populating database. This can take a while...')

        if self.is_populated():
            self.clear_database()

        start_index = 0
        processed_count = 0

        url = 'https://services.nvd.nist.gov/rest/json/cves/2.0'
        headers = {'apiKey': self.api_key} if self.api_key else None

        while True:
            params = {
                'startIndex': start_index,
                'noRejected': None
            }

            response = requests.get(url, headers=headers, params=params)
            if response.status_code == 200:
                data = response.json()
                if not data:
                    break

                for cve in data['vulnerabilities']:
                    cve['_id'] = cve['cve']['id']

                self.cve_col.insert_many(data['vulnerabilities'])
                processed_count += data['resultsPerPage']
                print(f"{processed_count}/{data['totalResults']} processed", end='\r')

                start_index += data['resultsPerPage']

                if start_index >= data['totalResults']:
                    # Finished populating database
                    self.cve_col.create_index({ 'cve.configurations.nodes.cpeMatch.criteria': pymongo.ASCENDING })
                    self.update_last_update_date(datetime.datetime.now())
                    print(f'{processed_count}/{data["totalResults"]} processed')
                    break

                time.sleep(2 if self.api_key else 6)

            else:
                print(f'Failed to fetch data from API: {response.status_code}')
                break

        print('Finished populating the CVE database')

    def update(self):
        print("Updating database...")

        if not self.is_populated():
            print("Database is not populated. Please populate the database first.")
            return

        print("Requesting changes since last update...")

        url = 'https://services.nvd.nist.gov/rest/json/cves/2.0'
        headers = {'apiKey': self.api_key} if self.api_key else None

        start_index = 0
        processed_count = 0

        last_update = self.get_last_update_date()

        while True:
            params = {
                'startIndex': start_index,
                'lastModStartDate': last_update.isoformat() if last_update else (datetime.datetime.now() - datetime.timedelta(weeks=1)).isoformat(),
                'lastModEndDate': datetime.datetime.now().isoformat(),
                'noRejected': None
            }

            response = requests.get(url, headers=headers, params=params)

            if response.status_code == 200:
                data = response.json()

                if not data:
                    break;

                if len(data['vulnerabilities']) == 0:
                    break;

                for cve in data['vulnerabilities']:
                    cve['_id'] = cve['cve']['id']
                    self.cve_col.update_one({'_id': cve['_id']}, {'$set': cve}, upsert=True)

                processed_count += data['resultsPerPage']
                print(f"{processed_count}/{data['totalResults']} changes processed", end='\r')

                start_index += data['resultsPerPage']

                if start_index >= data['totalResults']:
                    print(f'{processed_count}/{data["totalResults"]} changes processed')
                    break

                time.sleep(2 if self.api_key else 6)

            else:
                print(f'Failed to fetch data from API: {response.status_code}')
                break

        print('Requesting new CVEs since last update...')

        start_index = 0
        processed_count = 0

        while True:
            params = {
                'startIndex': start_index,
                'kevStartDate': last_update.isoformat() if last_update else (datetime.datetime.now() - datetime.timedelta(weeks=1)).isoformat(),
                'kevEndDate': datetime.datetime.now().isoformat(),
                'noRejected': None
            }

            response = requests.get(url, headers=headers, params=params)

            if response.status_code == 200:
                data = response.json()

                if not data:
                    break;

                if len(data['vulnerabilities']) == 0:
                    self.update_last_update_date(datetime.datetime.now())
                    break;

                for cve in data['vulnerabilities']:
                    cve['_id'] = cve['cve']['id']

                self.cve_col.insert_many(data['vulnerabilities'])

                processed_count += data['resultsPerPage']
                print(f"{processed_count}/{data['totalResults']} new CVEs processed", end='\r')

                start_index += data['resultsPerPage']

                if start_index >= data['totalResults']:
                    print(f'{processed_count}/{data["totalResults"]} new CVEs processed')
                    self.update_last_update_date(datetime.datetime.now())
                    break

                time.sleep(2 if self.api_key else 6)

            else:
                print(f'Failed to fetch data from API: {response.status_code}')
                break

        print('Finished updating the CVE database')

    def clear_database(self):
        self.cve_col.delete_many({})
        self.info_col.delete_many({})
        print('Database cleared')

    def is_populated(self) -> bool:
        return self.cve_col.estimated_document_count() > 0

    def get_last_update_date(self):
        last_update = self.info_col.find_one({}, sort=[('last_update', -1)])
        if last_update:
            return parser.parse(last_update['last_update'])
        return None

    def is_up_to_date(self) -> bool:
        last_update_date = self.get_last_update_date()
        if last_update_date:
            return last_update_date > datetime.datetime.now() - datetime.timedelta(weeks=1)
        return False

    def update_last_update_date(self, date: datetime.datetime):
        self.info_col.update_one({}, {'$set': {'last_update': date.isoformat()}}, upsert=True)

class HashDatabase(Database):
    def __init__(self, hash_query_url: str, api_key: str | None):
        super().__init__(hash_query_url)
        self.api_key: str | None = api_key
        self.hash_lookup: Hashlookup = Hashlookup()

    def query_database(self, package: str, hash: str) -> tuple[str, str, list[dict]]:
        return (package, hash, [self.hash_lookup.lookup(hash)])

    def connect(self) -> bool:
        return len(self.hash_lookup.info()) > 0