path: root/app/discover/scan_manager.py

###############################################################################
# Copyright (c) 2017 Koren Lev (Cisco Systems), Yaron Yogev (Cisco Systems)   #
# and others                                                                  #
#                                                                             #
# All rights reserved. This program and the accompanying materials            #
# are made available under the terms of the Apache License, Version 2.0       #
# which accompanies this distribution, and is available at                    #
# http://www.apache.org/licenses/LICENSE-2.0                                  #
###############################################################################
import argparse
import datetime

import time

import pymongo
from functools import partial

from discover.manager import Manager
from utils.constants import ScanStatus, EnvironmentFeatures
from utils.exceptions import ScanArgumentsError
from utils.inventory_mgr import InventoryMgr
from utils.logging.file_logger import FileLogger
from utils.mongo_access import MongoAccess
from discover.scan import ScanController


class ScanManager(Manager):

    DEFAULTS = {
        "mongo_config": "",
        "scans": "scans",
        "scheduled_scans": "scheduled_scans",
        "environments": "environments_config",
        "interval": 1,
        "loglevel": "INFO"
    }

    def __init__(self):
        self.args = self.get_args()
        super().__init__(log_directory=self.args.log_directory,
                         mongo_config_file=self.args.mongo_config)
        self.db_client = None
        self.environments_collection = None
        self.scans_collection = None
        self.scheduled_scans_collection = None

    @staticmethod
    def get_args():
        parser = argparse.ArgumentParser()
        parser.add_argument("-m", "--mongo_config", nargs="?", type=str,
                            default=ScanManager.DEFAULTS["mongo_config"],
                            help="Name of config file " +
                                 "with MongoDB server access details")
        parser.add_argument("-c", "--scans_collection", nargs="?", type=str,
                            default=ScanManager.DEFAULTS["scans"],
                            help="Scans collection to read from")
        parser.add_argument("-s", "--scheduled_scans_collection", nargs="?",
                            type=str,
                            default=ScanManager.DEFAULTS["scheduled_scans"],
                            help="Scans collection to read from")
        parser.add_argument("-e", "--environments_collection", nargs="?",
                            type=str,
                            default=ScanManager.DEFAULTS["environments"],
                            help="Environments collection to update "
                                 "after scans")
        parser.add_argument("-i", "--interval", nargs="?", type=float,
                            default=ScanManager.DEFAULTS["interval"],
                            help="Interval between collection polls"
                                 "(must be more than {} seconds)"
                                 .format(ScanManager.MIN_INTERVAL))
        parser.add_argument("-l", "--loglevel", nargs="?", type=str,
                            default=ScanManager.DEFAULTS["loglevel"],
                            help="Logging level \n(default: '{}')"
                                 .format(ScanManager.DEFAULTS["loglevel"]))
        parser.add_argument("-d", "--log_directory", nargs="?", type=str,
                            default=FileLogger.LOG_DIRECTORY,
                            help="File logger directory \n(default: '{}')"
                                 .format(FileLogger.LOG_DIRECTORY))
        args = parser.parse_args()
        return args

    def configure(self):
        self.db_client = MongoAccess()
        self.inv = InventoryMgr()
        self.inv.set_collections()
        self.scans_collection = self.db_client.db[self.args.scans_collection]
        self.scheduled_scans_collection = \
            self.db_client.db[self.args.scheduled_scans_collection]
        self.environments_collection = \
            self.db_client.db[self.args.environments_collection]
        self._update_document = \
            partial(MongoAccess.update_document, self.scans_collection)
        self.interval = max(self.MIN_INTERVAL, self.args.interval)
        self.log.set_loglevel(self.args.loglevel)

        self.log.info("Started ScanManager with following configuration:\n"
                      "Mongo config file path: {0.args.mongo_config}\n"
                      "Scans collection: {0.scans_collection.name}\n"
                      "Environments collection: "
                      "{0.environments_collection.name}\n"
                      "Polling interval: {0.interval} second(s)"
                      .format(self))

    def _build_scan_args(self, scan_request: dict):
        args = {
            'mongo_config': self.args.mongo_config
        }

        def set_arg(name_from: str, name_to: str = None):
            if name_to is None:
                name_to = name_from
            val = scan_request.get(name_from)
            if val:
                args[name_to] = val

        set_arg("object_id", "id")
        set_arg("log_level", "loglevel")
        set_arg("environment", "env")
        set_arg("scan_only_inventory", "inventory_only")
        set_arg("scan_only_links", "links_only")
        set_arg("scan_only_cliques", "cliques_only")
        set_arg("inventory")
        set_arg("clear")
        set_arg("clear_all")

        return args

    def _finalize_scan(self, scan_request: dict, status: ScanStatus,
                       scanned: bool):
        scan_request['status'] = status.value
        self._update_document(scan_request)
        # If no object id is present, it's a full env scan.
        # We need to update environments collection
        # to reflect the scan results.
        if not scan_request.get('id'):
            self.environments_collection\
                .update_one(filter={'name': scan_request.get('environment')},
                            update={'$set': {'scanned': scanned}})

    def _fail_scan(self, scan_request: dict):
        self._finalize_scan(scan_request, ScanStatus.FAILED, False)

    def _complete_scan(self, scan_request: dict, result_message: str):
        status = ScanStatus.COMPLETED if result_message == 'ok' \
            else ScanStatus.COMPLETED_WITH_ERRORS
        self._finalize_scan(scan_request, status, True)

    # PyCharm type checker can't reliably check types of document
    # noinspection PyTypeChecker
    def _clean_up(self):
        # Find and fail all running scans
        running_scans = list(self
                             .scans_collection
                             .find(filter={'status': ScanStatus.RUNNING.value}))
        self.scans_collection \
            .update_many(filter={'_id': {'$in': [scan['_id']
                                                 for scan
                                                 in running_scans]}},
                         update={'$set': {'status': ScanStatus.FAILED.value}})

        # Find all environments connected to failed full env scans
        env_scans = [scan['environment']
                     for scan in running_scans
                     if not scan.get('object_id')
                     and scan.get('environment')]

        # Set 'scanned' flag in those envs to false
        if env_scans:
            self.environments_collection\
                .update_many(filter={'name': {'$in': env_scans}},
                             update={'$set': {'scanned': False}})

    INTERVALS = {
        'YEARLY': datetime.timedelta(days=365.25),
        'MONTHLY': datetime.timedelta(days=365.25/12),
        'WEEKLY': datetime.timedelta(weeks=1),
        'DAILY': datetime.timedelta(days=1),
        'HOURLY': datetime.timedelta(hours=1)
    }

    def _submit_scan_request_for_schedule(self, scheduled_scan, interval, ts):
        scans = self.scans_collection
        new_scan = {
            'status': 'submitted',
            'log_level': scheduled_scan['log_level'],
            'clear': scheduled_scan['clear'],
            'scan_only_inventory': scheduled_scan['scan_only_inventory'],
            'scan_only_links': scheduled_scan['scan_only_links'],
            'scan_only_cliques': scheduled_scan['scan_only_cliques'],
            'submit_timestamp': ts,
            'interval': interval,
            'environment': scheduled_scan['environment'],
            'inventory': 'inventory'
        }
        scans.insert_one(new_scan)

    def _set_scheduled_requests_next_run(self, scheduled_scan, interval, ts):
        scheduled_scan['scheduled_timestamp'] = ts + self.INTERVALS[interval]
        doc_id = scheduled_scan.pop('_id')
        self.scheduled_scans_collection.update({'_id': doc_id}, scheduled_scan)

    def _prepare_scheduled_requests_for_interval(self, interval):
        now = datetime.datetime.utcnow()

        # first, submit a scan request where the scheduled time has come
        condition = {'$and': [
            {'freq': interval},
            {'scheduled_timestamp': {'$lte': now}}
        ]}
        matches = self.scheduled_scans_collection.find(condition) \
            .sort('scheduled_timestamp', pymongo.ASCENDING)
        for match in matches:
            self._submit_scan_request_for_schedule(match, interval, now)
            self._set_scheduled_requests_next_run(match, interval, now)

        # now set scheduled time where it was not set yet (new scheduled scans)
        condition = {'$and': [
            {'freq': interval},
            {'scheduled_timestamp': {'$exists': False}}
        ]}
        matches = self.scheduled_scans_collection.find(condition)
        for match in matches:
            self._set_scheduled_requests_next_run(match, interval, now)

    def _prepare_scheduled_requests(self):
        # see if any scheduled request is waiting to be submitted
        for interval in self.INTERVALS.keys():
            self._prepare_scheduled_requests_for_interval(interval)

    def do_action(self):
        self._clean_up()
        try:
            while True:
                self._prepare_scheduled_requests()

                # Find a pending request that is waiting the longest time
                results = self.scans_collection \
                    .find({'status': ScanStatus.PENDING.value,
                           'submit_timestamp': {'$ne': None}}) \
                    .sort("submit_timestamp", pymongo.ASCENDING) \
                    .limit(1)

                # If no scans are pending, sleep for some time
                if results.count() == 0:
                    time.sleep(self.interval)
                else:
                    scan_request = results[0]
                    env = scan_request.get('environment')
                    scan_feature = EnvironmentFeatures.SCANNING
                    if not self.inv.is_feature_supported(env, scan_feature):
                        self.log.error("Scanning is not supported for env '{}'"
                                       .format(scan_request.get('environment')))
                        self._fail_scan(scan_request)
                        continue

                    scan_request['start_timestamp'] = datetime.datetime.utcnow()
                    scan_request['status'] = ScanStatus.RUNNING.value
                    self._update_document(scan_request)

                    # Prepare scan arguments and run the scan with them
                    try:
                        scan_args = self._build_scan_args(scan_request)

                        self.log.info("Starting scan for '{}' environment"
                                      .format(scan_args.get('env')))
                        self.log.debug("Scan arguments: {}".format(scan_args))
                        result, message = ScanController().run(scan_args)
                    except ScanArgumentsError as e:
                        self.log.error("Scan request '{id}' "
                                       "has invalid arguments. "
                                       "Errors:\n{errors}"
                                       .format(id=scan_request['_id'],
                                               errors=e))
                        self._fail_scan(scan_request)
                    except Exception as e:
                        self.log.exception(e)
                        self.log.error("Scan request '{}' has failed."
                                       .format(scan_request['_id']))
                        self._fail_scan(scan_request)
                    else:
                        # Check is scan returned success
                        if not result:
                            self.log.error(message)
                            self.log.error("Scan request '{}' has failed."
                                           .format(scan_request['_id']))
                            self._fail_scan(scan_request)
                            continue

                        # update the status and timestamps.
                        self.log.info("Request '{}' has been scanned. ({})"
                                      .format(scan_request['_id'], message))
                        end_time = datetime.datetime.utcnow()
                        scan_request['end_timestamp'] = end_time
                        self._complete_scan(scan_request, message)
        finally:
            self._clean_up()


if __name__ == "__main__":
    ScanManager().run()