summaryrefslogtreecommitdiffstats
path: root/app/discover/scan_manager.py
diff options
context:
space:
mode:
Diffstat (limited to 'app/discover/scan_manager.py')
-rw-r--r--app/discover/scan_manager.py294
1 files changed, 294 insertions, 0 deletions
diff --git a/app/discover/scan_manager.py b/app/discover/scan_manager.py
new file mode 100644
index 0000000..b6ad782
--- /dev/null
+++ b/app/discover/scan_manager.py
@@ -0,0 +1,294 @@
+###############################################################################
+# 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
+
+ @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):
+ self._finalize_scan(scan_request, ScanStatus.COMPLETED, 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,
+ '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]
+ if not self.inv.is_feature_supported(scan_request.get('environment'),
+ EnvironmentFeatures.SCANNING):
+ 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']))
+ end_time = datetime.datetime.utcnow()
+ scan_request['end_timestamp'] = end_time
+ self._complete_scan(scan_request)
+ finally:
+ self._clean_up()
+
+
+if __name__ == "__main__":
+ ScanManager().run()