class TestCompare(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=MockDbInterface(),
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):
with pytest.raises(AttributeError):
self.compare_system.compare(['error'])
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.uid], 'dev_1')
self.assertEqual(result['device_name'][self.fw_two.uid], 'dev_2')
self.assertEqual(result['device_class'][self.fw_one.uid], 'Router')
self.assertEqual(result['vendor'][self.fw_one.uid], 'test_vendor')
self.assertEqual(result['version'][self.fw_one.uid], '0.1')
self.assertEqual(result['release_date'][self.fw_one.uid], '1970-01-01')
self.assertEqual(result['size'][self.fw_one.uid],
len(self.fw_one.binary))
self.assertEqual(result['virtual_file_path'][self.fw_one.uid],
[self.fw_one.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')
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
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