def _generate_file_tree(self, root_uid: str, uid: str, whitelist: List[str]) -> FileTreeNode:
root = FileTreeNode(None)
with ConnectTo(FrontEndDbInterface, self._config) as sc:
child_uids = [
child_uid
for child_uid in sc.get_specific_fields_of_db_entry(uid, {'files_included': 1})['files_included']
if whitelist is None or child_uid in whitelist
]
for node in sc.generate_file_tree_nodes_for_uid_list(child_uids, root_uid or uid, whitelist):
root.add_child_node(node)
return root
def _generate_exclusive_files_tree(self, exclusive_files, root_uid, uid):
root = FileTreeNode(None)
with ConnectTo(FrontEndDbInterface, self._config) as sc:
child_uids = sc.get_specific_fields_of_db_entry(
uid, {'files_included': 1})['files_included']
for child_uid in child_uids:
if not exclusive_files or child_uid in exclusive_files:
for node in sc.generate_file_tree_node(
child_uid, root_uid, whitelist=exclusive_files):
root.add_child_node(node)
return root
def get_node_for_virtual_file(self, current_virtual_path, fo_data,
root_uid, uid, whitelist):
node = FileTreeNode(uid=None,
root_uid=root_uid,
virtual=True,
name=current_virtual_path.pop(0))
for child_node in self.generate_file_tree_node(
uid,
root_uid,
current_virtual_path=current_virtual_path,
fo_data=fo_data,
whitelist=whitelist):
node.add_child_node(child_node)
return node
def _ajax_get_tree_children(self, uid, root_uid=None, compare_id=None):
exclusive_files = self._get_exclusive_files(compare_id, root_uid)
children = []
root = FileTreeNode(None)
with ConnectTo(FrontEndDbInterface, self._config) as sc:
child_uids = sc.get_specific_fields_of_db_entry(
uid, {'files_included': 1})['files_included']
for child_uid in child_uids:
if not exclusive_files or child_uid in exclusive_files:
for node in sc.generate_file_tree_node(
child_uid, root_uid, whitelist=exclusive_files):
root.add_child_node(node)
for child_node in root.get_list_of_child_nodes():
child = self._generate_jstree_node(child_node)
children.append(child)
return jsonify(children)
def generate_file_tree_level(self, uid, root_uid, whitelist=None, fo_data=None):
if fo_data is None:
fo_data = self.get_specific_fields_of_db_entry({'_id': uid}, VirtualPathFileTree.FO_DATA_FIELDS)
try:
for node in VirtualPathFileTree(root_uid, fo_data, whitelist).get_file_tree_nodes():
yield node
except (KeyError, TypeError): # the requested data is not in the DB aka the file has not been analyzed yet
yield FileTreeNode(uid, root_uid, not_analyzed=True, name='{uid} (not analyzed yet)'.format(uid=uid))
def get_node_for_real_file(self, current_virtual_path, fo_data, root_uid,
uid, whitelist):
mime_type = fo_data['processed_analysis'].get(
'file_type', {}).get('mime') or 'file-type-plugin/not-run-yet'
return FileTreeNode(uid,
root_uid=root_uid,
virtual=False,
name=self._get_file_name(current_virtual_path,
fo_data),
size=fo_data['size'],
mime_type=mime_type,
has_children=self._has_children(
fo_data, whitelist))
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))
def _create_node_from_virtual_path(self, uid, root_uid, current_virtual_path, fo_data, whitelist=None):
if len(current_virtual_path) > 1: # in the middle of a virtual file path
node = FileTreeNode(uid=None, root_uid=root_uid, virtual=True, name=current_virtual_path.pop(0))
for child_node in self.generate_file_tree_node(uid, root_uid, current_virtual_path=current_virtual_path, fo_data=fo_data, whitelist=whitelist):
node.add_child_node(child_node)
else: # at the end of a virtual path aka a 'real' file
if whitelist:
has_children = any(f in fo_data['files_included'] for f in whitelist)
else:
has_children = fo_data['files_included'] != []
mime_type = fo_data['processed_analysis']['file_type']['mime'] if 'file_type' in fo_data['processed_analysis'] else 'file-type-plugin/not-run-yet'
node = FileTreeNode(uid, root_uid=root_uid, virtual=False, name=fo_data['file_name'], size=fo_data['size'], mime_type=mime_type, has_children=has_children)
return node
def test_node_merging(self):
parent_node = FileTreeNode('123',
virtual=False,
name='parent',
size=1,
mime_type='somefile')
child_node_folder_1 = FileTreeNode(None, virtual=True, name='folder')
child_node_folder_2 = FileTreeNode(None, virtual=True, name='folder')
child_node_file_1 = FileTreeNode('abc', virtual=False, name='file_1')
child_node_file_2 = FileTreeNode('def', virtual=False, name='file_2')
child_node_folder_1.add_child_node(child_node_file_1)
child_node_folder_2.add_child_node(child_node_file_2)
parent_node.add_child_node(child_node_folder_1)
parent_node.add_child_node(child_node_folder_2)
assert parent_node.has_children
assert len(parent_node.get_list_of_child_nodes()) == 1
assert list(
parent_node.children.keys()) == [child_node_folder_1.get_id()]
assert child_node_folder_1 in parent_node
assert child_node_folder_2 in parent_node
assert len(parent_node.children[
child_node_folder_1.get_id()].get_list_of_child_nodes()) == 2
folder_id = child_node_folder_1.get_id()
assert child_node_file_1 in parent_node.children[folder_id]
assert child_node_file_2 in parent_node.children[folder_id]
def test_node_creation(self):
parent_node = FileTreeNode('123',
virtual=False,
name='parent',
size=1,
mime_type='somefile')
child_node = FileTreeNode('456',
root_uid='123',
virtual=True,
name='child')
parent_node.add_child_node(child_node)
assert parent_node.uid == '123'
assert parent_node.root_uid is None
assert child_node.root_uid == '123'
assert not parent_node.virtual
assert parent_node.size == 1
assert parent_node.type == 'somefile'
assert parent_node.has_children
assert parent_node.get_list_of_child_nodes() == [child_node]
assert parent_node.get_id() == ('parent', False)
assert list(parent_node.children.keys()) == [child_node.get_id()]
assert parent_node.get_names_of_children() == [child_node.name]
assert not child_node.has_children
assert child_node in parent_node
assert 'Node ' in parent_node.__repr__()
assert parent_node != child_node
assert parent_node.print_tree() is None
def test_node_merging(self):
parent_node = FileTreeNode('123',
virtual=False,
name='parent',
size=1,
mime_type='somefile')
child_node_folder_1 = FileTreeNode(None, virtual=True, name='folder')
child_node_folder_2 = FileTreeNode(None, virtual=True, name='folder')
child_node_file_1 = FileTreeNode('abc', virtual=False, name='file_1')
child_node_file_2 = FileTreeNode('def', virtual=False, name='file_2')
child_node_folder_1.add_child_node(child_node_file_1)
child_node_folder_2.add_child_node(child_node_file_2)
parent_node.add_child_node(child_node_folder_1)
parent_node.add_child_node(child_node_folder_2)
self.assertTrue(parent_node.has_children)
self.assertEqual(len(parent_node.get_list_of_child_nodes()), 1)
self.assertEqual(list(parent_node.children.keys()),
[child_node_folder_1.get_id()])
self.assertTrue(child_node_folder_1 in parent_node)
self.assertTrue(child_node_folder_2 in parent_node)
self.assertEqual(
len(parent_node.children[
child_node_folder_1.get_id()].get_list_of_child_nodes()), 2)
folder_id = child_node_folder_1.get_id()
self.assertTrue(child_node_file_1 in parent_node.children[folder_id])
self.assertTrue(child_node_file_2 in parent_node.children[folder_id])
def test_node_creation(self):
parent_node = FileTreeNode('123',
virtual=False,
name='parent',
size=1,
mime_type='somefile')
child_node = FileTreeNode('456',
root_uid='123',
virtual=True,
name='child')
parent_node.add_child_node(child_node)
self.assertEqual(parent_node.uid, '123')
self.assertEqual(parent_node.root_uid, None)
self.assertEqual(child_node.root_uid, '123')
self.assertFalse(parent_node.virtual)
self.assertEqual(parent_node.size, 1)
self.assertEqual(parent_node.type, 'somefile')
self.assertTrue(parent_node.has_children)
self.assertEqual(parent_node.get_list_of_child_nodes(), [child_node])
self.assertEqual(parent_node.get_id(), ('parent', False))
self.assertEqual(list(parent_node.children.keys()),
[child_node.get_id()])
self.assertEqual(parent_node.get_names_of_children(),
[child_node.name])
self.assertFalse(child_node.has_children)
self.assertTrue(child_node in parent_node)
assert 'Node ' in parent_node.__repr__()
self.assertNotEqual(parent_node, child_node)
assert parent_node.print_tree() is None
def test_node_creation(self):
parent_node = FileTreeNode("123", False, "parent", 1, "somefile")
child_node = FileTreeNode("456", True, "child")
parent_node.add_child_node(child_node)
self.assertEqual(parent_node.uid, "123")
self.assertFalse(parent_node.virtual)
self.assertEqual(parent_node.size, 1)
self.assertEqual(parent_node.type, "somefile")
self.assertTrue(parent_node.has_children)
self.assertEqual(parent_node.get_list_of_child_nodes(), [child_node])
self.assertEqual(parent_node.get_id(), ("parent", False))
self.assertEqual(list(parent_node.children.keys()), [child_node.get_id()])
self.assertEqual(parent_node.get_names_of_children(), [child_node.name])
self.assertFalse(child_node.has_children)
self.assertTrue(child_node in parent_node)
self.assertTrue(parent_node, FileTreeNode("123", False, "parent", 1, "somefile"))
def test_node_merging(self):
parent_node = FileTreeNode("123", False, "parent", 1, "somefile")
child_node_folder_1 = FileTreeNode(None, True, "folder")
child_node_folder_2 = FileTreeNode(None, True, "folder")
child_node_file_1 = FileTreeNode("abc", False, "file_1")
child_node_file_2 = FileTreeNode("def", False, "file_2")
child_node_folder_1.add_child_node(child_node_file_1)
child_node_folder_2.add_child_node(child_node_file_2)
parent_node.add_child_node(child_node_folder_1)
parent_node.add_child_node(child_node_folder_2)
self.assertTrue(parent_node.has_children)
self.assertEqual(len(parent_node.get_list_of_child_nodes()), 1)
self.assertEqual(list(parent_node.children.keys()), [child_node_folder_1.get_id()])
self.assertTrue(child_node_folder_1 in parent_node)
self.assertTrue(child_node_folder_2 in parent_node)
self.assertEqual(len(parent_node.children[child_node_folder_1.get_id()].get_list_of_child_nodes()), 2)
folder_id = child_node_folder_1.get_id()
self.assertTrue(child_node_file_1 in parent_node.children[folder_id])
self.assertTrue(child_node_file_2 in parent_node.children[folder_id])
