easy-vuls

#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
Preform a scan of host using vuls.
"""
import argparse
import logging
import os
import sqlite3
import sys
from datetime import date, datetime
from shutil import rmtree
from subprocess import run


class VulsScanException(Exception):
    """
    Class for exceptions which are may occur due to issues beyond control of
    this developer.
    """


class DatabaseDoesNotExist(VulsScanException):
    """
    This will be raised if there is an attempt to access the database before
    it has been created.
    """


class DockerManager(object):
    REQUIRED_DIRECTORIES = []

    def __init__(self):
        self.ensure_directories()
        self.environs = ()
        self.volumes = ()

    @property
    def path_prefix(self):
        return os.path.abspath(sys.argv[0])

    def _path(self, path):
        """
        Return a path relative to this script.
        e.g. /home/user/script.path/foo
        """
        return "{}.{}".format(self.path_prefix, path)

    def ensure_directories(self):
        """
        Ensure all the required directories exist.
        """
        directories = [self._path(d) for d in self.REQUIRED_DIRECTORIES]
        for directory in directories:
            if not os.path.isdir(directory):
                if os.path.exists(directory):
                    raise ValueError("{} is not a directory".format(directory))
                os.makedirs(directory)

    def run(self, log, args, image='vuls/vuls', publish=()):
        """
        Run a container with base+additional arguments and fix volume owners.
        """
        log_directory = self._path('logs/{}'.format(log))
        self._run(log, args, image=image, publish=publish)

        # fix file ownership
        volumes = []
        for source, _ in self.volumes + ((log_directory, ''), ):
            volumes.append('--volume={0}:/tmp/{0}'.format(source))
        # copy owner:group from the script
        stats = os.stat(os.path.abspath(sys.argv[0]))
        chown = 'chown -R {}:{} /tmp/'.format(
                    stats.st_uid, stats.st_gid,
                    )
        run([
                'docker', 'run', '--rm', '--interactive',
                '--entrypoint=/bin/sh',
                ] + volumes + [
                'vuls/vuls', '-c', chown,
        ])

    def _run(self, log, args, image, publish):
        """
        Run a container with base+additional arguments.
        """
        docker_args = [
                'run', '--rm', '--interactive', '--tty',
                '--volume={}:/var/log/vuls'.format(
                    self._path('logs/{}'.format(log))
                    ),
                '--volume={}:/vuls'.format(self._path('data')),
                '--volume=/etc/localtime:/etc/localtime:ro',
        ]
        for source, destination in self.volumes:
            docker_args.append('--volume={}:{}'.format(source, destination))
        for key, value in self.environs:
            docker_args.append('--env={}={}'.format(key, value))
        if publish:
            listen_address, host_port, container_port = publish
            docker_args.append('--publish={}:{}:{}'.format(
                listen_address,
                host_port,
                container_port,
            ))

        run(['docker'] + docker_args + [image] + args)


class NVD(DockerManager):
    """
    National Vulnerability Database
    """
    REQUIRED_DIRECTORIES = [
        'data',
        'logs',
    ]

    def __init__(self):
        super().__init__()
        self.database_path = self._path('data/cve.sqlite3')

    @property
    def database_initialised(self):
        """
        Indicate if the database has been initialised.
        """
        return (os.path.isfile(self.database_path)
                and os.path.getsize(self.database_path) > 0)

    def get_new_cursor(self):
        """
        DRY way of getting a database cursor.
        """
        if not self.database_initialised:
            raise DatabaseDoesNotExist()

        if not hasattr(self, '_database_connection'):
            self._database_connection = sqlite3.connect(
                    self.database_path,
                    detect_types=sqlite3.PARSE_COLNAMES,
                    )

        return self._database_connection.cursor()

    def get_cve_stats(self):
        """
        Return the number of CVEs and the last time any was updated.
        """
        total, last_modified = self.get_new_cursor().execute('''
                SELECT
                    Count(*),
                    Max(updated_at) as "[timestamp]"
                FROM nvds
                ''').fetchone()
        return (total, last_modified)

    def update(self):
        """
        Refresh the NVD copy. Return a boolean indicating if any updates where
        made, and the number of CVEs added.

        Avoid running this in the minutes that a new year starts to avoid
        missing the last few updates of the year. Probably incredibly unlikely.
        """
        if not self.database_initialised:
            old_total = 0
            update_from = datetime(2002, 1, 1)
            with open(self.database_path, 'wb'):
                pass
        else:
            old_total, update_from = self.get_cve_stats()

        self.pull_year_cve_data(range(update_from.year, date.today().year + 1))
        new_total, updated_to = self.get_cve_stats()

        return (update_from != updated_to, new_total - old_total)

    def pull_year_cve_data(self, years):
        """
        Pull NVD updates and store them in the persistent sqlite3 database.
        """
        args = [
            'fetchnvd', '-years',
        ] + [str(year) for year in years]
        self.run('fetchnvd', args, image='vuls/go-cve-dictionary')

    def delete(self):
        """
        Remove the database.
        """
        try:
            os.remove(self.database_path)
        except FileNotFoundError:
            pass


class Vuls(DockerManager):
    REQUIRED_DIRECTORIES = [
            'data',
            'data/keys'
            'logs',
            'results',
    ]

    def __init__(self, config_path, debug=False):
        super().__init__()
        self.config_path = os.path.abspath(config_path)
        self.results_directory = self._path('results')
        self.debug = debug

        self.volumes += (
            (self.results_directory, self.results_directory),
            (self.config_path, '/vuls/config.toml'),
        )

        ssh_auth_sock = os.environ.get('SSH_AUTH_SOCK')
        if ssh_auth_sock is None:
            logging.warn("No SSH_AUTH_SOCK found. See ssh-agent")
        else:
            self.environs += (
                ('SSH_AUTH_SOCK', '/root/SSH_AUTH_SOCK'),
            )
            self.volumes += (
                (ssh_auth_sock, '/root/SSH_AUTH_SOCK'),
            )

    def prepare(self):
        """
        Prepare hosts for scanning if necissary.
        """
        self.run('prepare', ['prepare'] + (['--debug'] if self.debug else []))

    def scan(self):
        """
        Scan the target machines with the provided config.
        """
        nvd = NVD()
        if not nvd.database_initialised:
            raise DatabaseDoesNotExist()

        args = [
                'scan',
                '--results-dir={}'.format(self.results_directory),
                '-cve-dictionary-dbpath=/vuls/cve.sqlite3',
                '-report-json',
        ] + (['--debug'] if self.debug else [])
        self.run('scan', args)


class Results(DockerManager):
    REQUIRED_DIRECTORIES = [
        'results',
    ]

    def __init__(self):
        super().__init__()
        self.volumes += (
            (self._path('results'), '/vuls/results'),
        )

    def browse_results(self):
        """
        Start up the ncurses based user interface to the scan results.
        """
        self.run('tui', ['tui'])

    def serve(self, listen_address, listen_port):
        """
        Start up VulsRepo to serve the results.
        """
        self.run(
            'vulsrepo',
            [],
            image='vuls/vulsrepo',
            publish=(listen_address, listen_port, 80),
        )

    def delete(self):
        """
        Delete the recorded results.
        """
        results_dir = self._path('results')
        for entry in os.scandir(results_dir):
            if entry.is_dir(follow_symlinks=False):
                rmtree(entry.path)
            else:
                os.remove(entry.path)


class VulsScannerInterface(object):
    """
    CLI for VulsScanner
    """

    @classmethod
    def run(cls, args):
        """
        Take args, pop the command, and create an instance to handle the run.
        """
        args = cls.parse_args(args)
        command = args.command
        del args.command
        try:
            cls(command, args)
        except DatabaseDoesNotExist:
            logging.error(
                    "Unable to preform operation as the database does not "
                    "yet exist. Run 'vuls-scan database update' first",
                    )

    def __init__(self, command, args):
        """
        Dispatch the arguments to the approriate command handler.
        """
        {
            'database': self.database,
            'scan': self.scan,
            'results': self.results,
        }[command](args)

    def database(self, args):
        nvd = NVD()
        if args.action == 'update':
            nvd.update()
        elif args.action == 'stats':
            entries, last_update = nvd.get_cve_stats()
            print("There are currently {} NVD entries.".format(entries))
            print("The last time any changes were retrieved was {}".format(last_update))
        elif args.action == 'delete':
            nvd.delete()

    def scan(self, args):
        """
        Run a scan with the given config
        """
        vuls = Vuls(args.config, debug=args.debug)
        if args.action == 'prepare':
            vuls.prepare()
        else:
            vuls.scan()

    def results(self, args):
        """
        Show the results of previous scans in some way.
        """
        results = Results()
        if args.action == 'browse':
            results.browse_results()
        elif args.action == 'serve':
            print("About to serve results at {}:{}/vulsrepo/".format(
                    args.address, args.port,
            ))
            results.serve(args.address, args.port)
        elif args.action == 'delete':
            results.delete()

    @classmethod
    def parse_args(cls, args):
        parser = argparse.ArgumentParser(
                description="Scan for unsecure packages using vuls.",
        )
        commands = parser.add_subparsers(dest='command', help="command help")
        commands.required = True  # see https://bugs.python.org/issue9253

        # maybe database update/delete/stats
        database = commands.add_parser(
                'database',
                help="Commands relating to the database",
                )
        database_action = database.add_subparsers(dest='action')
        database_action.required = True
        database_action.add_parser(
                'update',
                help="Retrieve CVE data and store it for future scans"
        )
        database_action.add_parser(
                'stats',
                help="Show stats for the current CVE data",
        )
        database_action.add_parser(
                'delete',
                help="Remove the CVE database entierly",
        )

        scan = commands.add_parser(
                'scan',
                help="Perform a scan on target hosts",
        )
        scan.add_argument(
                '--config',
                required=True,
                help="The config file to use for vuls",
        )
        scan.add_argument(
                '--debug',
                action='store_true',
                help="Make vuls log debug information",
        )
        scan_action = scan.add_subparsers(
                dest='action',
        )
        scan_action.add_parser(
                'prepare',
                help="Prepare hosts for scanning if necissary",
        )

        results = commands.add_parser(
                'results',
                help="Commands to handle the results",
                )
        results_action = results.add_subparsers(dest='action')
        results_action.required = True
        results_action.add_parser(
                'delete',
                help="Delete all of the stored results",
        )
        results_action.add_parser(
                'browse',
                help="Browse the results in the terminal"
        )
        results_serve = results_action.add_parser(
                'serve',
                help="Serve the results for viewing",
        )
        results_serve.add_argument(
                '--port',
                default='8080',
                help="Port to serve content on",
        )
        results_serve.add_argument(
                '--address',
                default='127.0.0.1',
                help="Address to listen on"
        )

        return parser.parse_args(args)


if __name__ == '__main__':
    logging.basicConfig(
        format='[%(asctime)s] %(levelname)s: %(message)s',
        level=logging.INFO,
        )
    VulsScannerInterface.run(None)