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)
async def test_run_after_non_exists(self, cfg): self.thread.NAME = 'test_name' cfg['portdetection.test_name.scan_enabled'] = True cfg['portdetection.test_name.run_after'] = ['test_scan'] cfg['portdetection.expiration_period'] = '7d' self.thread.run = MagicMock(return_value=Future()) self.thread.run.return_value.set_result(True) task = MagicMock(NAME='test_scan_other') self.thread.aucote.async_task_manager = AsyncTaskManager() self.thread.aucote.async_task_manager.add_crontab_task( task=task, cron='', event=MagicMock()) await self.thread() self.assertFalse(task.run_asap.called)
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()
def get_app(self): self.aucote = MagicMock(unfinished_tasks=4) self.context = ScanContext(aucote=self.aucote, scanner=None) self.tasks = AsyncTaskManager() self.tasks._is_running = False tasks = [ { 'port': 45, 'id': 1, 'ip': '127.0.0.1' }, { 'port': 56, 'id': 2, 'ip': '127.0.0.2' }, { 'port': 67, 'id': 3, 'ip': '127.0.0.3' }, ] worker_tasks = [ { 'port': 78, 'id': 4, 'ip': '127.0.0.4' }, { 'port': 89, 'id': 5, 'ip': '127.0.0.5' }, ] for task in tasks: self.tasks.add_task( PortTask(context=self.context, port=Port( node=Node(node_id=task['id'], ip=ipaddress.ip_address(task['ip'])), number=task['port'], transport_protocol=TransportProtocol.TCP, ), exploits=[])) self.tasks._task_workers = { number: PortTask(context=self.context, port=Port( node=Node(node_id=task['id'], ip=ipaddress.ip_address(task['ip'])), number=task['port'], transport_protocol=TransportProtocol.TCP, ), exploits=[]) for number, task in enumerate(worker_tasks) } class test_function: def __str__(self): return 'test_function' def __call__(self, *args, **kwargs): pass self.tasks.add_crontab_task(test_function(), '0 0 0 0 0') self.tasks._task_workers.update({2: None, 3: None, 4: None}) self.aucote.async_task_managers = {TaskManagerType.SCANNER: self.tasks} self.app = Application([(r"/api/v1/tasks", TasksHandler, { 'aucote': self.aucote })]) return self.app
class TasksHandlerTest(AsyncHTTPTestCase): def setUp(self): super(TasksHandlerTest, self).setUp() def get_app(self): self.aucote = MagicMock(unfinished_tasks=4) self.context = ScanContext(aucote=self.aucote, scanner=None) self.tasks = AsyncTaskManager() self.tasks._is_running = False tasks = [ { 'port': 45, 'id': 1, 'ip': '127.0.0.1' }, { 'port': 56, 'id': 2, 'ip': '127.0.0.2' }, { 'port': 67, 'id': 3, 'ip': '127.0.0.3' }, ] worker_tasks = [ { 'port': 78, 'id': 4, 'ip': '127.0.0.4' }, { 'port': 89, 'id': 5, 'ip': '127.0.0.5' }, ] for task in tasks: self.tasks.add_task( PortTask(context=self.context, port=Port( node=Node(node_id=task['id'], ip=ipaddress.ip_address(task['ip'])), number=task['port'], transport_protocol=TransportProtocol.TCP, ), exploits=[])) self.tasks._task_workers = { number: PortTask(context=self.context, port=Port( node=Node(node_id=task['id'], ip=ipaddress.ip_address(task['ip'])), number=task['port'], transport_protocol=TransportProtocol.TCP, ), exploits=[]) for number, task in enumerate(worker_tasks) } class test_function: def __str__(self): return 'test_function' def __call__(self, *args, **kwargs): pass self.tasks.add_crontab_task(test_function(), '0 0 0 0 0') self.tasks._task_workers.update({2: None, 3: None, 4: None}) self.aucote.async_task_managers = {TaskManagerType.SCANNER: self.tasks} self.app = Application([(r"/api/v1/tasks", TasksHandler, { 'aucote': self.aucote })]) return self.app def test_tasks(self): expected = { 'unfinished_tasks': 4, 'scanner': { 'unfinished_tasks': 3, 'queue': [ '[+] [None] PortTask [on 127.0.0.1:45]', '[+] [None] PortTask [on 127.0.0.2:56]', '[+] [None] PortTask [on 127.0.0.3:67]' ], 'workers': { 'count': 5, 'jobs': { '0': '[+] [None] PortTask [on 127.0.0.4:78]', '1': '[+] [None] PortTask [on 127.0.0.5:89]', '2': None, '3': None, '4': None } }, 'cron_tasks': [{ 'name': 'test_function', 'cron': '0 0 0 0 0', 'is_running': False }] } } response = self.fetch('/api/v1/tasks', method='GET') self.assertEqual(response.code, 200) self.assertEqual(response.headers['Content-Type'], "application/json; charset=UTF-8") self.assertEqual(json.loads(response.body.decode()), expected)