cve_db.py

#!/usr/bin/env python3

import argparse
import datetime
import logging
import time
import socket

import db
import nvd


logger = logging.getLogger('MAIN')


def parse_host(host: str) -> str:
    try:
        socket.getaddrinfo(host, None)
    except socket.gaierror:
        raise ValueError
    return str(host)


def parse_port(port: str) -> int:
    p = int(port)
    if 0 < p < 65535:
        return p
    raise ValueError


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description='Synchronize a local CVE mongodb database with remote NVD data feeds')

    conn_info = parser.add_argument_group('MongoDB connection information')
    conn_info.add_argument('--host', type=parse_host, required=True, help='MongoDB host')
    conn_info.add_argument('--port', type=parse_port, default=27017, help='MongoDB port')
    conn_info.add_argument('--db', type=str, required=True, help='MongoDB database name')
    conn_info.add_argument('--username', type=str, default=None, help='MongoDB login username')
    conn_info.add_argument('--password', type=str, default=None, help='MongoDB login password')
    conn_info.add_argument('--authentication-source', type=str, default='admin', help='MongoDB authentication source')
    conn_info.add_argument('--clear-existing', action='store_true', default=False, help='MongoDB clear existing'
                                                                                           ' collections')

    feed_info = parser.add_argument_group('Feed parser settings')
    feed_info.add_argument('--nvd', metavar='URL', type=str, default='https://nvd.nist.gov/vuln/data-feeds',
                           help='NVD data feed source URL')

    return parser.parse_args()


def setup_logging():
    logging.basicConfig(format='[%(asctime)-15s] [%(name)s] [%(levelname)s] %(message)s', level=logging.INFO,
                        datefmt='%Y-%m-%d %H:%M:%S')


def live_sources_synchronizable(live: nvd.SourceList[nvd.Source]) -> bool:
    for src in live:
        if db.feed_unknown(src):
            return False
        modified_delta = db.feed_modified_delta(src)
        if modified_delta >= datetime.timedelta(days=8):
            logging.warning(f'The live feed \'{src.name}\' has not been synchronized for more than 8 days.')
            return False
    return True


def sources_up_to_date(sources: nvd.SourceList[nvd.Source]) -> bool:
    for src in sources:
        modified_delta = db.feed_modified_delta(src)
        if modified_delta != datetime.timedelta(days=0):
            return False
    return True


def rebuild(live: nvd.SourceList[nvd.Source], archive: nvd.SourceList[nvd.Source]):
    logging.warning(f'Rebuilding local database to mitigate desynchronization with NVD')
    archive.pull_data()
    live.pull_data()
    for sources in (archive, live):
        for src in sources:
            db.feed_migrate_data(src)
            db.feed_migrate_meta(src)
    archive.flush_data()
    live.flush_data()


def update(live: nvd.SourceList[nvd.Source]):
    if sources_up_to_date(live):
        logger.info(f'Live sources are up to date')
        return
    logger.info(f'Updating database using live source deltas')
    live.pull_data()
    for src in live:
        db.feed_migrate_data(src)
        db.feed_migrate_meta(src)
    live.flush_data()


def update_loop(aggregator: nvd.SourceAggregator):
    while True:
        aggregator.update_sources()
        aggregator.live_sources.pull_meta()

        if live_sources_synchronizable(live=aggregator.live_sources):
            update(live=aggregator.live_sources)
        else:
            aggregator.archive_sources.pull_meta()
            rebuild(live=aggregator.live_sources, archive=aggregator.archive_sources)

        delta = datetime.timedelta(hours=2)
        now = datetime.datetime.now()
        logger.info(f'Next update at: {now + delta} (2 hours from now)')
        time.sleep(delta.total_seconds())


def update_loop_failsafe(aggregator: nvd.SourceAggregator):
    backoff = 30
    while True:
        try:
            update_loop(aggregator)
        except KeyboardInterrupt as k:
            raise k
        except Exception as e:
            logger.exception(f'Caught exception:\n')
            logger.error(f'Backing off for {backoff} seconds')
            time.sleep(backoff)
            logger.error(f'Resuming')
            backoff = min(backoff + 30, 300)


def main():
    setup_logging()
    args = parse_args()

    with db.connect(**vars(args)):
        if args.clear_existing:
            db.clear()
        aggregator = nvd.SourceAggregator(root=args.nvd)
        try:
            update_loop_failsafe(aggregator)
        except KeyboardInterrupt:
            logger.info('Keyboard interrupt received, exiting')
        aggregator.close()


if __name__ == "__main__":
    main()