def test_subtree(self): subtree_copy = Tree(self.tree.subtree("jane"), deep=True) self.assertEqual(subtree_copy.parent("jane") is None, True) subtree_copy["jane"].tag = "Sweeti" self.assertEqual(self.tree["jane"].tag == "Jane", True) self.assertEqual(subtree_copy.level("diane"), 1) self.assertEqual(subtree_copy.level("jane"), 0) self.assertEqual(self.tree.level("jane"), 1)
def get_lca(T: tl.Tree, x: int, y: int) -> int: # First, get to the same level while T.level(x) > T.level(y): x = T.parent(x).identifier while T.level(x) < T.level(y): y = T.parent(y).identifier # Then, increment both until it's the same node. while x != y: x = T.parent(x).identifier y = T.parent(y).identifier # now, this is the LCA. return x
def test_shallow_paste(self): t1 = Tree() n1 = t1.create_node(identifier='A') t2 = Tree() n2 = t2.create_node(identifier='B') t3 = Tree() n3 = t3.create_node(identifier='C') t1.paste(n1.identifier, t2) self.assertEqual(t1.to_dict(), {'A': {'children': ['B']}}) t1.paste(n1.identifier, t3) self.assertEqual(t1.to_dict(), {'A': {'children': ['B', 'C']}}) self.assertEqual(t1.level(n1.identifier), 0) self.assertEqual(t1.level(n2.identifier), 1) self.assertEqual(t1.level(n3.identifier), 1)
class StepParse: def __init__(self): pass def load_step(self, step_filename): self.nauo_lines = [] self.prod_def_lines = [] self.prod_def_form_lines = [] self.prod_lines = [] self.filename = os.path.splitext(step_filename)[0] line_hold = '' line_type = '' # Find all search lines with open(step_filename) as f: for line in f: # TH: read pointer of lines as they are read, so if the file has text wrap it will notice and add it to the following lines index = re.search("#(.*)=", line) if index: # TH: if not none then it is the start of a line so read it # want to hold line until it has checked next line # if next line is a new indexed line then save previous line if line_hold: if line_type == 'nauo': self.nauo_lines.append(line_hold) elif line_type == 'prod_def': self.prod_def_lines.append(line_hold) elif line_type == 'prod_def_form': self.prod_def_form_lines.append(line_hold) elif line_type == 'prod': self.prod_lines.append(line_hold) line_hold = '' line_type = '' prev_index = True # TH remember previous line had an index if 'NEXT_ASSEMBLY_USAGE_OCCURRENCE' in line: line_hold = line.rstrip() line_type = 'nauo' elif ('PRODUCT_DEFINITION ' in line or 'PRODUCT_DEFINITION(' in line): line_hold = line.rstrip() line_type = 'prod_def' elif 'PRODUCT_DEFINITION_FORMATION' in line: line_hold = line.rstrip() line_type = 'prod_def_form' elif ('PRODUCT ' in line or 'PRODUCT(' in line): line_hold = line.rstrip() line_type = 'prod' else: prev_index = False #TH: if end of file and previous line was held if 'ENDSEC;' in line: if line_hold: if line_type == 'nauo': self.nauo_lines.append(line_hold) elif line_type == 'prod_def': self.prod_def_lines.append(line_hold) elif line_type == 'prod_def_form': self.prod_def_form_lines.append(line_hold) elif line_type == 'prod': self.prod_lines.append(line_hold) line_hold = '' line_type = '' else: #TH: if not end of file line_hold = line_hold + line.rstrip() self.nauo_refs = [] self.prod_def_refs = [] self.prod_def_form_refs = [] self.prod_refs = [] # TH: added 'replace(","," ").' to replace ',' with a space to make the spilt easier if there are not spaces inbetween the words' # Find all (# hashed) line references and product names # TH: it might be worth finding a different way of extracting data we do want rather than fixes to get rid of the data we don't for j, el_ in enumerate(self.nauo_lines): self.nauo_refs.append([ el.rstrip(',') for el in el_.replace(",", " ").replace("=", " ").split() if el.startswith('#') ]) for j, el_ in enumerate(self.prod_def_lines): self.prod_def_refs.append([ el.rstrip(',') for el in el_.replace(",", " ").replace("=", " ").split() if el.startswith('#') ]) for j, el_ in enumerate(self.prod_def_form_lines): self.prod_def_form_refs.append([ el.rstrip(',') for el in el_.replace(",", " ").replace("=", " ").split() if el.startswith('#') ]) for j, el_ in enumerate(self.prod_lines): self.prod_refs.append([ el.strip(',') for el in el_.replace(",", " ").replace( "(", " ").replace("=", " ").split() if el.startswith('#') ]) self.prod_refs[j].append(el_.split("'")[1]) # Get first two items in each sublist (as third is shape ref) # # First item is 'PRODUCT_DEFINITION' ref # Second item is 'PRODUCT_DEFINITION_FORMATION <etc>' ref self.prod_all_refs = [el[:2] for el in self.prod_def_refs] # Match up all references down to level of product name for j, el_ in enumerate(self.prod_all_refs): # Add 'PRODUCT_DEFINITION' ref for i, el in enumerate(self.prod_def_form_refs): if el[0] == el_[1]: el_.append(el[1]) break # Add names from 'PRODUCT_DEFINITION' lines for i, el in enumerate(self.prod_refs): if el[0] == el_[2]: el_.append(el[2]) break # Find all parent and child relationships (3rd and 2nd item in each sublist) self.parent_refs = [el[1] for el in self.nauo_refs] self.child_refs = [el[2] for el in self.nauo_refs] # Find distinct parts and assemblies via set operations; returns list, so no repetition of items self.all_type_refs = set(self.child_refs) | set(self.parent_refs) self.ass_type_refs = set(self.parent_refs) self.part_type_refs = set(self.child_refs) - set(self.parent_refs) #TH: find root node self.root_type_refs = set(self.parent_refs) - set(self.child_refs) # Create simple parts dictionary (ref + label) self.part_dict = {el[0]: el[3] for el in self.prod_all_refs} # self.part_dict_inv = {el[3]:el[0] for el in self.prod_all_refs} def show_values(self): # TH: basic testing, if needed these could be spilt up print(self.nauo_lines) print(self.prod_def_lines) print(self.prod_def_form_lines) print(self.prod_lines) print(self.nauo_refs) print(self.prod_def_refs) print(self.prod_def_form_refs) print(self.prod_refs) # HR: "create_dict" replaced by list comprehension elsewhere # # def create_dict(self): # # # TH: links nauo number with a name and creates dict # self.part_dict = {} # for part in self.all_type_refs: # for sublist in self.prod_def_refs: # if sublist[0] == part: # prod_loc = '#' + re.findall('\d+',sublist[1])[0] # pass # for sublist in self.prod_def_form_refs: # if sublist[0] == prod_loc: # prod_loc = '#' + str(re.findall('\d+',sublist[1])[0]) # pass # for sublist in self.prod_refs: # if sublist[0] == prod_loc: # part_name = sublist[2] # # self.part_dict[part] = part_name def create_tree(self): #TH: create tree diagram in newick format #TH: find root node self.tree = Tree() #TH: check if there are any parts to make a tree from, if not don't bother if self.part_dict == {}: return root_node_ref = list(self.root_type_refs)[0] # HR added part reference as data for later use self.tree.create_node(self.part_dict[root_node_ref], 0, data={'ref': root_node_ref}) #TH: created root node now fill in next layer #TH: create dict for tree, as each node needs a unique name i = [0] # Iterates through nodes self.tree_dict = {} self.tree_dict[i[0]] = root_node_ref def tree_next_layer(self, parent): root_node = self.tree_dict[i[0]] for line in self.nauo_refs: if line[1] == root_node: i[0] += 1 self.tree_dict[i[0]] = str(line[2]) # HR added part reference as data for later use self.tree.create_node(self.part_dict[line[2]], i[0], parent=parent, data={'ref': str(line[2])}) tree_next_layer(self, i[0]) tree_next_layer(self, 0) self.appended = False self.get_levels() def get_levels(self): # Initialise dict and get first level (leaves) self.levels = {} self.levels_set_p = set() self.levels_set_a = set() self.leaf_ids = [el.identifier for el in self.tree.leaves()] self.all_ids = [el for el in self.tree.nodes] self.non_leaf_ids = set(self.all_ids) - set(self.leaf_ids) self.part_level = 1 def do_level(self, tree_level): # Get all nodes within this level node_ids = [ el for el in self.tree.nodes if self.tree.level(el) == tree_level ] for el in node_ids: # If leaf, then n_p = 1 and n_a = 1 if el in self.leaf_ids: self.levels[el] = {} self.levels[el]['n_p'] = self.part_level self.levels[el]['n_a'] = self.part_level # If assembly, then get all children and sum all parts + assemblies else: # Get all children of node and sum levels child_ids = self.tree.is_branch(el) child_sum_p = 0 child_sum_a = 0 for el_ in child_ids: child_sum_p += self.levels[el_]['n_p'] child_sum_a += self.levels[el_]['n_a'] self.levels[el] = {} self.levels[el]['n_p'] = child_sum_p self.levels[el]['n_a'] = child_sum_a + 1 self.levels_set_p.add(child_sum_p) self.levels_set_a.add(child_sum_a + 1) # Go up through tree levels and populate lattice level dict for i in range(self.tree.depth(), -1, -1): do_level(self, i) self.create_lattice() self.levels_p_sorted = sorted(list(self.levels_set_p)) self.levels_a_sorted = sorted(list(self.levels_set_a)) # Function to return dictionary of item IDs for each lattice level def get_levels_inv(list_in, key): #Initialise levels_inv = {} levels_inv[self.part_level] = [] for el in list_in: levels_inv[el] = [] for k, v in self.levels.items(): levels_inv[v[key]].append(k) return levels_inv self.levels_p_inv = get_levels_inv(self.levels_p_sorted, 'n_p') self.levels_a_inv = get_levels_inv(self.levels_a_sorted, 'n_a') def get_all_children(self, id_): ancestors = [el.identifier for el in self.tree.children(id_)] parents = ancestors while parents: children = [] for parent in parents: children = [el.identifier for el in self.tree.children(parent)] ancestors.extend(children) parents = children return ancestors def create_lattice(self): # Create lattice self.g = nx.DiGraph() self.default_colour = 'r' # Get root node and set parent to -1 to maintain data type of "parent" # Set position to top/middle node_id = self.tree.root label_text = self.tree.get_node(node_id).tag self.g.add_node(node_id, parent=-1, label=label_text, colour=self.default_colour) # Do nodes from treelib "nodes" dictionary for key in self.tree.nodes: # Exclude root if key != self.tree.root: parent_id = self.tree.parent(key).identifier label_text = self.tree.get_node(key).tag # Node IDs same as for tree self.g.add_node(key, parent=parent_id, label=label_text, colour=self.default_colour) # Do edges from nodes for key in self.tree.nodes: # Exclude root if key != self.tree.root: parent_id = self.tree.parent(key).identifier self.g.add_edge(key, parent_id) # Escape if only one node # HR 6/3/20 QUICK BUG FIX: SINGLE-NODE TREE DOES NOT PLOT # IMPROVE LATER; SHOULD BE PART OF A GENERAL METHOD if self.tree.size() == 1: id_ = [el.identifier for el in self.tree.leaves()] self.g.nodes[id_[-1]]['pos'] = (0, 0) return # Get set of parents of leaf nodes leaf_parents = set( [self.tree.parent(el).identifier for el in self.leaf_ids]) # For each leaf_parent, set position of leaf nodes sequentially i = 0 no_leaves = len(self.tree.leaves()) for el in leaf_parents: for el_ in self.tree.is_branch(el): child_ids = [el.identifier for el in self.tree.leaves()] if el_ in child_ids: self.g.nodes[el_]['pos'] = ((i / (no_leaves)), 1) i += 1 # To set plot positions of nodes from lattice levels # --- # Traverse upwards from leaves for el in sorted(list(self.levels_set_a)): # Get all nodes at that level node_ids = [k for k, v in self.levels.items() if v['n_a'] == el] # Get all positions of children of that node # and set position as mean value of them for el_ in node_ids: child_ids = self.tree.is_branch(el_) pos_sum = 0 for el__ in child_ids: pos_ = self.g.nodes[el__]['pos'][0] pos_sum += pos_ pos_sum = pos_sum / len(child_ids) self.g.nodes[el_]['pos'] = (pos_sum, el) def print_tree(self): try: self.tree.show() except: self.create_tree() self.tree.show() def tree_to_json(self, save_to_file=False, filename='file', path=''): #TH: return json format tree, can also save to file if self.tree.size() != 0: data = self.tree.to_json() j = json.loads(data) if save_to_file == True: if path: file_path = os.path.join(path, filename) else: file_path = filename with open(file_path + '.json', 'w') as outfile: json.dump(j, outfile) return data else: print("no tree to print") return
"Select extra knowledge to load mirror infomation") if filename is not None: with open(filename, 'r') as load_f: load_dict = json.load(load_f) mirror = get_mirror(root.startEA) procs = get_all_procs() tree.create_node(fname, hex(root.startEA), data=Xref_node(fname, hex(root.startEA), XType.code, mirror)) if mirror == fname: add_xrefs(root.startEA, XType.code) Message("Reference Tree:\n\n") tree.show(line_type="ascii-em", idhidden=False, data_property='mirror') Message("Unique references:\n") for node in tree.all_nodes_itr(): if type(node.data.mirror) is str: print node.identifier #hierarchical output for level in range(1, tree.depth()): Message("\nLevel %d: %d\n" % (level, tree.size(level))) for node in tree.all_nodes(): if tree.level(node.identifier) == level and type( node.data.mirror) is str: print node.identifier Message("\n%d subroutines in routine %s need transplanting.\n" % (Xref_node.xrefTrans - 1, fname)) conn.close() else: Warning("No function found at location %x" % here())
def level_size(tree: treelib.Tree, level=1): return len( [x for x in tree.all_nodes() if tree.level(x.identifier) == level])
class IOTeqDBBuilder(): def __init__(self, configFile): with open(configFile, 'r') as f: jsonOutput = json.load(f) self.databaseTags = jsonOutput["database"]["tags"] self.IOTEQ_TAG_BYTE_SIZE = 44 self.tree = Tree() self.constPtrChar = [] self.constPtrTree = [] self.tagList = [] self.dataPtr = [] self.persistentPtr = [] self.currentTagAddress = 0 def totalNumberOfTags(self): return len(self.tagList) def addTag(self, tag): self.tagList.append(tag) def addNameToCharPtr(self, name): charList = list(name) for char in charList: self.constPtrChar.append(ord(char)) self.constPtrChar.append(0) # escape char for string def setRoot(self, rootName): rootTag = IOTeqTag(rootName, 0x00, len(rootName)) self.tree.create_node(rootName, rootName, None, rootTag) self.tagList.append(rootTag) self.addNameToCharPtr(rootName) def addValueToDataPtr(self, tag): datatype = tag["datatype"] if (datatype == "Number"): value = tag["value"] if ("numtype" in tag["config"]): if (tag["config"]["numtype"] == "float"): ba = bytearray(struct.pack("<f", value)) else: ba = bytearray(struct.pack("<L", value)) for b in ba: self.dataPtr.append(hex(b)) elif (datatype == "Text"): value = tag["value"] for char in value: self.dataPtr.append(hex(ord(char))) # for i in range(len(self.dataPtr), 40): self.dataPtr.append(hex(0)) def addValueToPersistentPtr(self, tag): datatype = tag["datatype"] if (datatype == "Number"): value = tag["value"] if ("numtype" in tag["config"]): if (tag["config"]["numtype"] == "float"): ba = bytearray(struct.pack("<f", value)) else: ba = bytearray(struct.pack("<L", value)) for b in ba: self.persistentPtr.append(hex(b)) elif (datatype == "Text"): value = tag["value"] for char in value: self.persistentPtr.append(hex(ord(char))) # for i in range(len(self.dataPtr), 40): self.persistentPtr.append(hex(0)) def createTree(self, tags, parent=None): for tag in tags: for i in range(tags[tag]["arraydim"]): # Naming function for tags with dimensions larger than 1 if (tags[tag]["arraydim"] > 1): tagName = tag + "[" + str(i) + "]" else: tagName = tag # Add tag name in hex format to a list charIndex = len(self.constPtrChar) self.addNameToCharPtr(tagName) # Create tag an add to tagList newTag = IOTeqTag(tagName, charIndex, len(tagName) + 1) # plus 1 for \0 # Set the tags valueSize based on datatype # Numbers are 4 bytes (float) and Text are 40 characters long total if (tags[tag]["datatype"] != "Folder"): if (tags[tag]["datatype"] == "Number"): newTag.valueSize = 4 elif (tags[tag]["datatype"] == "Text"): newTag.valueSize = len(tags[tag]["value"]) if ("persistent" in tags[tag]["config"]): if (tags[tag]["config"]["persistent"] == True): newTag.persistentValuePtr = len(self.persistentPtr) self.addValueToPersistentPtr(tags[tag]) newTag.isPersistent = 1 # Adding default value to dataPtr list newTag.valuePtr = len(self.dataPtr) self.addValueToDataPtr(tags[tag]) # Add num type to tag object if ("numtype" in tags[tag]["config"]): if (tags[tag]["config"]["numtype"] == "float"): newTag.numType = "float" else: newTag.numType = "integer" # Set default value of tag newTag.value = tags[tag]['value'] # Adding tag to tree and tag list self.tree.create_node(tagName, tagName, parent, newTag) self.tagList.append(newTag) # Recursion for tags that have children if (tags[tag]["datatype"] == "Folder"): self.createTree(tags[tag]["children"], tagName) def setTagAddresses(self): for level in range(self.tree.depth() + 1): for node in dict( filter(lambda elem: self.tree.level(elem[0]) == level, self.tree.nodes.items())): tag = self.tree.get_node(node).data tag.address = self.currentTagAddress self.currentTagAddress += self.IOTEQ_TAG_BYTE_SIZE def setTagParentChildrenPtrs(self): for node in self.tree.nodes: # tagIndex = None for tag in self.tagList: if (tag.tagName == node): tagIndex = self.tagList.index(tag) # Get IOTeq Tag treeNode = self.tree.get_node(node) self.tagList[tagIndex] = treeNode.data # If node has a parent, i.e. not root if (treeNode.is_root() != True): # Get IOTeq Parent Tag parentTag = self.tree.parent(node).data # Set Parent Ptr of current tag self.tagList[tagIndex].parentPtr = parentTag.address if (parentTag.tagName != "tags"): self.tagList[tagIndex].parentTag = parentTag # If node has children if (self.tree.get_node(node).is_leaf() != True): childrenNodes = self.tree.children(node) childrenNodes.sort(key=lambda x: x.data.address) self.tagList[tagIndex].childPtr = childrenNodes[ 0].data.address self.tagList[tagIndex].numOfChildren = len( childrenNodes) if (tag.tagName != "tags"): for i in range(0, len(childrenNodes)): tagIndex = self.tagList.index( childrenNodes[i].data) if (i == 0): self.tagList[ tagIndex].nextSibling = childrenNodes[ i + 1].data.address elif (i == len(childrenNodes) - 1): self.tagList[ tagIndex].prevSibling = childrenNodes[ i - 1].data.address else: self.tagList[ tagIndex].nextSibling = childrenNodes[ i + 1].data.address self.tagList[ tagIndex].prevSibling = childrenNodes[ i - 1].data.address def createConstPtrTree(self): sortedTags = sorted(self.tagList, key=lambda x: x.address, reverse=False) for tag in sortedTags: self.constPtrTree.extend(tag.getStruct()) def build(self): self.setRoot("tags") self.createTree(ioteqDBBuilder.databaseTags, "tags") self.setTagAddresses() self.setTagParentChildrenPtrs() self.createConstPtrTree()