def setUp(self, cfg):
super(ScannerTest, self).setUp()
self.cfg = {
'portdetection': {
'scan_type': 'LIVE',
'live_scan': {
'min_time_gap': 0,
},
'_internal': {
'tools_cron': '* * * * *'
}
},
'topdis': {
'api': {
'host': '',
'port': ''
},
}
}
cfg._cfg = self.cfg
self.aucote = MagicMock(storage=MagicMock())
atm_stop_future = Future()
self.atm_stop = MagicMock()
atm_stop_future.set_result(self.atm_stop)
self.aucote.async_task_manager.stop.return_value = atm_stop_future
self.thread = Scanner(aucote=self.aucote)
self.thread._init()
self.thread._cron_tasks = {1: MagicMock(), 2: MagicMock()}
self.task_manager = AsyncTaskManager.instance()
self.task_manager.run_tasks = {'_run_tools': False, '_scan': False}
def setUp(self):
super(TestAsyncTaskManager, self).setUp()
self.task_1 = MagicMock()
self.task_2 = MagicMock()
AsyncTaskManager._instances = {}
self.task_manager = AsyncTaskManager.instance(parallel_tasks=1)
self.task_manager._shutdown_condition = MagicMock()
self.task_manager._stop_condition = MagicMock()
self.task_manager._cron_tasks['task_1'] = self.task_1
self.task_manager._cron_tasks['task_2'] = self.task_2
self.task_manager._task_workers = {0: None}
def __init__(self, exploits, kudu_queue, tools_config):
self.exploits = exploits
self._kudu_queue = kudu_queue
signal.signal(signal.SIGINT, self.signal_handler)
signal.signal(signal.SIGTERM, self.signal_handler)
signal.signal(signal.SIGHUP, self.graceful_stop)
self.load_tools(tools_config)
self._storage = Storage(conn_string=cfg['storage.db'], nodes_limit=cfg['storage.max_nodes_query'])
self.ioloop = IOLoop.current()
self.topdis = Topdis(cfg['topdis.api.host'], cfg['topdis.api.port'], cfg['topdis.api.base'])
self.async_task_managers = {
TaskManagerType.SCANNER: AsyncTaskManager.instance(name=TaskManagerType.SCANNER.value,
parallel_tasks=10),
TaskManagerType.REGULAR: AsyncTaskManager.instance(name=TaskManagerType.REGULAR.value,
parallel_tasks=cfg['service.scans.parallel_tasks']),
TaskManagerType.QUICK: AsyncTaskManager.instance(name=TaskManagerType.QUICK.value, parallel_tasks=30)
}
for task_manager in self.async_task_managers.values():
throttling_consumer = ThrottlingConsumer(manager=task_manager)
self.ioloop.add_callback(partial(cfg.add_rabbit_consumer, throttling_consumer))
self.ioloop.add_callback(throttling_consumer.consume)
self.web_server = WebServer(self, cfg['service.api.v1.host'], cfg['service.api.v1.port'],
path=cfg['service.api.path'])
self._tftp_thread = TFTPThread(cfg['tftp.host'], cfg['tftp.port'], timeout=cfg['tftp.timeout'],
min_port=cfg['tftp.min_port'], max_port=cfg['tftp.max_port'])
self._storage_thread = StorageThread(storage=self._storage)
self.scanners = []
if cfg.toucan:
for task_manager in self.async_task_managers.values():
cfg.toucan_monitor.register_toucan_key(key='throttling.rate', add_prefix=False, default=1,
callback=task_manager.change_throttling_toucan)
def setUp(self, cfg):
super(ScanAsyncTaskTest, self).setUp()
self.cfg = {
'portdetection': {
'test_name': {
'scan_type': 'LIVE',
'live_scan': {
'min_time_gap': 0,
},
'scan_devices': {
'snmp': True,
'host': True
}
},
'_internal': {
'tools_cron': '* * * * *'
},
'expiration_period': '365d'
},
'topdis': {
'api': {
'host': '',
'port': ''
},
}
}
cfg._cfg = self.cfg
self.aucote = MagicMock()
atm_stop_future = Future()
self.atm_stop = MagicMock()
atm_stop_future.set_result(self.atm_stop)
self.aucote.async_task_manager.stop.return_value = atm_stop_future
self.thread = ScanAsyncTask(aucote=self.aucote)
self.thread.NAME = 'test_name'
self.thread._cron_tasks = {1: MagicMock(), 2: MagicMock()}
self.task_manager = AsyncTaskManager.instance()
self.task_manager.run_tasks = {'_run_tools': False, '_scan': False}
def tearDown(self):
AsyncTaskManager.instance().clear()