def setUp(self): self.config = get_config_for_testing() self.fw_one = create_test_firmware(device_name='dev_1', all_files_included_set=True) self.fw_one.processed_analysis['file_hashes'] = {'ssdeep': get_ssdeep(self.fw_one.binary)} self.fw_two = create_test_firmware(device_name='dev_2', bin_path='container/test.7z', all_files_included_set=True) self.fw_two.processed_analysis['file_hashes'] = {'ssdeep': get_ssdeep(self.fw_two.binary)} self.compare_system = Compare(db_interface=MockDbInterface(), config=self.config)
class Test_Compare(unittest.TestCase): def setUp(self): self.config = get_config_for_testing() self.fw_one = create_test_firmware(device_name='dev_1', all_files_included_set=True) self.fw_one.processed_analysis['file_hashes'] = { 'ssdeep': get_ssdeep(self.fw_one.binary) } self.fw_two = create_test_firmware(device_name='dev_2', bin_path='container/test.7z', all_files_included_set=True) self.fw_two.processed_analysis['file_hashes'] = { 'ssdeep': get_ssdeep(self.fw_two.binary) } self.compare_system = Compare(db_interface=mock_db_interface(), config=self.config) def tearDown(self): gc.collect() def test_compare_objects(self): result = self.compare_system.compare_objects( [self.fw_one, self.fw_two]) self.assertIsInstance(result, dict, 'Result is not a dict') self.assertIn('general', result, 'general part is missing') self.assertIsInstance(result['general'], dict, 'general part is not a dict') self.assertIn('plugins', result, 'plugin part is missing') self.assertIsInstance(result['plugins'], dict, 'plugins part is not a dict') def test_compare_error_none_existing_fo(self): result = self.compare_system.compare(['error']) self.assertIsInstance(result, Exception, 'result has wrong type') def test_create_general_section_dict(self): result = self.compare_system._create_general_section_dict( [self.fw_one, self.fw_two]) self.assertIsInstance(result, dict, 'result is not a dict') self.assertEqual(result['device_name'][self.fw_one.get_uid()], 'dev_1') self.assertEqual(result['device_name'][self.fw_two.get_uid()], 'dev_2') self.assertEqual(result['device_class'][self.fw_one.get_uid()], 'Router') self.assertEqual(result['vendor'][self.fw_one.get_uid()], 'test_vendor') self.assertEqual(result['version'][self.fw_one.get_uid()], '0.1') self.assertEqual(result['release_date'][self.fw_one.get_uid()], '1970-01-01') self.assertEqual(result['size'][self.fw_one.get_uid()], len(self.fw_one.binary)) self.assertEqual(result['virtual_file_path'][self.fw_one.get_uid()], [self.fw_one.get_uid()]) def test_plugin_system(self): self.assertGreater(len(self.compare_system.compare_plugins), 0, 'no compare plugin found') self.assertIn('File_Coverage', self.compare_system.compare_plugins, 'File Coverage module not found')
def __init__(self, config=None, db_interface=None, testing=False, callback=None): self.config = config self.db_interface = db_interface if db_interface else CompareDbInterface(config=config) self.stop_condition = Value('i', 1) self.in_queue = Queue() self.callback = callback self.compare_module = Compare(config=self.config, db_interface=self.db_interface) self.worker = ExceptionSafeProcess(target=self._compare_scheduler_main) if not testing: self.start()
class CompareScheduler: ''' This module handles all request regarding compares ''' def __init__(self, config=None, db_interface=None, testing=False, callback=None): self.config = config self.db_interface = db_interface if db_interface else CompareDbInterface( config=config) self.stop_condition = Value('i', 1) self.in_queue = Queue() self.callback = callback self.compare_module = Compare(config=self.config, db_interface=self.db_interface) self.worker = ExceptionSafeProcess(target=self._compare_scheduler_main) if not testing: self.start() def start(self): self.stop_condition.value = 0 self.worker.start() logging.info('Compare Scheduler online...') def shutdown(self): ''' shutdown the scheduler ''' logging.debug('Shutting down...') if getattr(self.db_interface, 'shutdown', False): self.db_interface.shutdown() if self.stop_condition.value == 0: self.stop_condition.value = 1 self.worker.join() self.in_queue.close() logging.info('Compare Scheduler offline') def add_task(self, compare_task): compare_id, redo = compare_task try: self.db_interface.check_objects_exist(compare_id) except FactCompareException as exception: return exception.get_message( ) # FIXME: return value gets ignored by backend intercom logging.debug('Schedule for compare: {}'.format(compare_id)) self.in_queue.put((compare_id, redo)) return None def _compare_scheduler_main(self): compares_done = set() while self.stop_condition.value == 0: self._compare_single_run(compares_done) logging.debug('Compare Thread terminated') def _compare_single_run(self, compares_done): try: compare_id, redo = self.in_queue.get( timeout=float(self.config['ExpertSettings']['block_delay'])) except Empty: pass else: if self._decide_whether_to_process(compare_id, redo, compares_done): if redo: self.db_interface.delete_old_compare_result(compare_id) compares_done.add(compare_id) self._process_compare(compare_id) if self.callback: self.callback() def _process_compare(self, compare_id): result = self.compare_module.compare( convert_compare_id_to_list(compare_id)) if isinstance(result, dict): self.db_interface.add_compare_result(result) else: logging.error(result) @staticmethod def _decide_whether_to_process(uid, redo, compares_done): return redo or uid not in compares_done def check_exceptions(self): return_value = False if self.worker.exception: logging.error("{}Worker Exception Found!!{}".format( bcolors.FAIL, bcolors.ENDC)) logging.error(self.worker.exception[1]) if self.config.getboolean('ExpertSettings', 'throw_exceptions'): return_value = True terminate_process_and_childs(self.worker) return return_value
def init_compare(hotspot_one, hotspot_two): """ Initialize comparison of hotspots """ path_for_one = "static/maps/" + hotspot_one + "/" + get_saved_csv( hotspot_one)[0] path_for_two = "static/maps/" + hotspot_two + "/" + get_saved_csv( hotspot_two)[0] csv_one = pandas.read_csv(path_for_one, usecols=[ "Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday" ], encoding="utf-8").values csv_two = pandas.read_csv(path_for_two, usecols=[ "Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday" ], encoding="utf-8").values compare = Compare(csv_one, csv_two) compare.calculate_overlap() compare.calculate_percentage() compare.calculate_jaccard() return { "data": compare.get_overlap(), "all_percentage": compare.get_percentage(), "jaccard": compare.get_jaccard(), # "pai": round(pai.pai_index(csv_one, csv_two), 3) # "z_max": use_max, # "z_min": use_min }
class CompareScheduler: ''' This module handles all request regarding compares ''' def __init__(self, config=None, db_interface=None, testing=False, callback=None): self.config = config self.db_interface = db_interface if db_interface else CompareDbInterface( config=config) self.stop_condition = Value('i', 1) self.in_queue = Queue() self.callback = callback self.compare_module = Compare(config=self.config, db_interface=self.db_interface) self.worker = ExceptionSafeProcess(target=self._compare_scheduler_main) if not testing: self.start() def start(self): self.stop_condition.value = 0 self.worker.start() logging.info('Compare Scheduler online...') def shutdown(self): ''' shutdown the scheduler ''' logging.debug('Shutting down...') if getattr(self.db_interface, 'shutdown', False): self.db_interface.shutdown() if self.stop_condition.value == 0: self.stop_condition.value = 1 self.worker.join() self.in_queue.close() logging.info('Compare Scheduler offline') def add_task(self, compare_task): compare_id, redo = compare_task try: self.db_interface.check_objects_exist(compare_id) except FactCompareException as exception: return exception.get_message( ) # FIXME: return value gets ignored by backend intercom logging.debug(f'Schedule for compare: {compare_id}') self.in_queue.put((compare_id, redo)) return None def _compare_scheduler_main(self): compares_done = set() while self.stop_condition.value == 0: self._compare_single_run(compares_done) logging.debug('Compare Thread terminated') def _compare_single_run(self, compares_done): try: compare_id, redo = self.in_queue.get( timeout=float(self.config['ExpertSettings']['block_delay'])) except Empty: pass else: if self._decide_whether_to_process(compare_id, redo, compares_done): if redo: self.db_interface.delete_old_compare_result(compare_id) compares_done.add(compare_id) self._process_compare(compare_id) if self.callback: self.callback() def _process_compare(self, compare_id): try: self.db_interface.add_compare_result( self.compare_module.compare( convert_compare_id_to_list(compare_id))) except Exception: # pylint: disable=broad-except logging.error(f'Fatal error in compare process for {compare_id}', exc_info=True) @staticmethod def _decide_whether_to_process(uid, redo, compares_done): return redo or uid not in compares_done def check_exceptions(self): processes_to_check = [self.worker] shutdown = check_worker_exceptions(processes_to_check, 'Compare', self.config, self._compare_scheduler_main) if not shutdown and new_worker_was_started( new_process=processes_to_check[0], old_process=self.worker): self.worker = processes_to_check.pop() return shutdown