class TestFsOrganizer(unittest.TestCase):
def setUp(self):
self.ds_tmp_dir = TemporaryDirectory(prefix='fact_tests_')
config = ConfigParser()
config.add_section('data_storage')
config.set('data_storage', 'firmware_file_storage_directory', self.ds_tmp_dir.name)
self.fs_organzier = FSOrganizer(config)
def tearDown(self):
self.ds_tmp_dir.cleanup()
gc.collect()
def check_file_presence_and_content(self, file_path, file_binary):
self.assertTrue(os.path.exists(file_path), 'file exists')
self.assertEqual(get_binary_from_file(file_path), file_binary, 'correct content')
def test_generate_path(self):
test_binary = b'abcde'
file_object = FileObject(test_binary)
file_path = self.fs_organzier.generate_path(file_object)
# file path should be 'DATA_DIR/UID_PEFIX/UID'
self.assertEqual(file_path, '{}/36/36bbe50ed96841d10443bcb670d6554f0a34b761be67ec9c4a8ad2c0c44ca42c_5'.format(self.ds_tmp_dir.name), 'generate file path')
def test_store_and_delete_file(self):
test_binary = b'abcde'
file_object = FileObject(test_binary)
self.fs_organzier.store_file(file_object)
self.check_file_presence_and_content('{}/36/36bbe50ed96841d10443bcb670d6554f0a34b761be67ec9c4a8ad2c0c44ca42c_5'.format(self.ds_tmp_dir.name), b'abcde')
self.assertEqual(file_object.file_path, '{}/36/36bbe50ed96841d10443bcb670d6554f0a34b761be67ec9c4a8ad2c0c44ca42c_5'.format(self.ds_tmp_dir.name), 'wrong file path set in file object')
self.fs_organzier.delete_file(file_object.uid)
self.assertFalse(os.path.exists(file_object.file_path), 'file not deleted')
class InterComBackEndReAnalyzeTask(InterComListener):
CONNECTION_TYPE = 're_analyze_task'
def additional_setup(self, config=None):
self.fs_organizer = FSOrganizer(config=config)
def post_processing(self, task, task_id):
task.file_path = self.fs_organizer.generate_path(task)
task.create_binary_from_path()
return task
class Unpacker(UnpackBase):
def __init__(self, config=None, worker_id=None, db_interface=None):
super().__init__(config=config, worker_id=worker_id)
self.file_storage_system = FSOrganizer(config=self.config)
self.db_interface = db_interface
def unpack(self, current_fo: FileObject):
'''
Recursively extract all objects included in current_fo and add them to current_fo.files_included
'''
logging.debug('[worker {}] Extracting {}: Depth: {}'.format(self.worker_id, current_fo.uid, current_fo.depth))
if current_fo.depth >= self.config.getint('unpack', 'max_depth'):
logging.warning('{} is not extracted since depth limit ({}) is reached'.format(current_fo.uid, self.config.get('unpack', 'max_depth')))
self._store_unpacking_depth_skip_info(current_fo)
return []
tmp_dir = TemporaryDirectory(prefix='fact_unpack_', dir=get_temp_dir_path(self.config))
file_path = self._generate_local_file_path(current_fo)
extracted_files = self.extract_files_from_file(file_path, tmp_dir.name)
extracted_file_objects = self.generate_and_store_file_objects(extracted_files, tmp_dir.name, current_fo)
extracted_file_objects = self.remove_duplicates(extracted_file_objects, current_fo)
self.add_included_files_to_object(extracted_file_objects, current_fo)
# set meta data
current_fo.processed_analysis['unpacker'] = json.loads(Path(tmp_dir.name, 'reports', 'meta.json').read_text())
self.cleanup(tmp_dir)
return extracted_file_objects
@staticmethod
def _store_unpacking_depth_skip_info(file_object: FileObject):
file_object.processed_analysis['unpacker'] = {
'plugin_used': 'None', 'number_of_unpacked_files': 0,
'info': 'Unpacking stopped because maximum unpacking depth was reached',
}
tag_dict = {'unpacker': {'depth reached': {'value': 'unpacking depth reached', 'color': TagColor.ORANGE, 'propagate': False}}}
file_object.analysis_tags.update(tag_dict)
def cleanup(self, tmp_dir):
try:
tmp_dir.cleanup()
except OSError as error:
logging.error('[worker {}] Could not CleanUp tmp_dir: {} - {}'.format(self.worker_id, type(error), str(error)))
@staticmethod
def add_included_files_to_object(included_file_objects, root_file_object):
for item in included_file_objects:
root_file_object.add_included_file(item)
def generate_and_store_file_objects(self, file_paths: List[Path], extractor_dir: str, parent: FileObject):
extracted_files = {}
for item in file_paths:
if not file_is_empty(item):
current_file = FileObject(file_path=str(item))
current_virtual_path = join_virtual_path(
get_base_of_virtual_path(parent.get_virtual_file_paths()[parent.get_root_uid()][0]),
parent.uid, get_object_path_excluding_fact_dirs(make_unicode_string(str(item)), str(Path(extractor_dir, 'files')))
)
current_file.temporary_data['parent_fo_type'] = get_file_type_from_path(parent.file_path)['mime']
if current_file.uid in extracted_files: # the same file is extracted multiple times from one archive
extracted_files[current_file.uid].virtual_file_path[parent.get_root_uid()].append(current_virtual_path)
else:
self.db_interface.set_unpacking_lock(current_file.uid)
self.file_storage_system.store_file(current_file)
current_file.virtual_file_path = {parent.get_root_uid(): [current_virtual_path]}
current_file.parent_firmware_uids.add(parent.get_root_uid())
extracted_files[current_file.uid] = current_file
return extracted_files
@staticmethod
def remove_duplicates(extracted_fo_dict, parent_fo):
if parent_fo.uid in extracted_fo_dict:
del extracted_fo_dict[parent_fo.uid]
return make_list_from_dict(extracted_fo_dict)
def _generate_local_file_path(self, file_object: FileObject):
if not Path(file_object.file_path).exists():
local_path = self.file_storage_system.generate_path(file_object.uid)
return local_path
return file_object.file_path