def test_get_partial_virtual_path(self):
virtual_path = {'abc': ['|abc|def|ghi|folder_1/folder_2/file']}
self.assertEqual(get_partial_virtual_path(virtual_path, 'abc'),
{'abc': ['|abc|def|ghi|folder_1/folder_2/file']})
self.assertEqual(get_partial_virtual_path(virtual_path, 'ghi'),
{'ghi': ['|ghi|folder_1/folder_2/file']})
self.assertEqual(get_partial_virtual_path(virtual_path, 'xyz'),
{'xyz': ['|xyz|']})
def test_get_partial_virtual_path(self):
virtual_path = {"abc": ["|abc|def|ghi|folder_1/folder_2/file"]}
self.assertEqual(get_partial_virtual_path(virtual_path, "abc"),
{'abc': ['|abc|def|ghi|folder_1/folder_2/file']})
self.assertEqual(get_partial_virtual_path(virtual_path, "ghi"),
{'ghi': ['|ghi|folder_1/folder_2/file']})
self.assertEqual(get_partial_virtual_path(virtual_path, "xyz"),
{'xyz': ['|xyz|']})
def generate_file_tree_node(self, uid, root_uid, current_virtual_path=None, fo_data=None, whitelist=None):
required_fields = {'virtual_file_path': 1, 'files_included': 1, 'file_name': 1, 'size': 1, 'processed_analysis.file_type.mime': 1, '_id': 1}
if fo_data is None:
fo_data = self.get_specific_fields_of_db_entry({'_id': uid}, required_fields)
try:
if root_uid not in fo_data['virtual_file_path']: # file tree for a file object (instead of a firmware)
fo_data['virtual_file_path'] = get_partial_virtual_path(fo_data['virtual_file_path'], root_uid)
if current_virtual_path is None:
for entry in fo_data['virtual_file_path'][root_uid]: # the same file may occur several times with different virtual paths
current_virtual_path = entry.split('/')[1:]
yield self._create_node_from_virtual_path(uid, root_uid, current_virtual_path, fo_data, whitelist)
else:
yield self._create_node_from_virtual_path(uid, root_uid, current_virtual_path, fo_data, whitelist)
except Exception: # the requested data is not present in the DB aka the file has not been analyzed yet
yield FileTreeNode(uid=uid, root_uid=root_uid, not_analyzed=True, name='{} (not analyzed yet)'.format(uid))
