def test_merge_source(self): """ The node it visits match the dest id :return: """ # # A relation from relation_source_id to relation_dest_id # relation_source_id = "source object id" relation_dest_id = "dest object id" source_node = TreeNode() source_node.id = relation_source_id des_node = TreeNode() des_node.id = relation_dest_id # This source node is the root of the tree tree = MultiRootTree(source_node) # The dest node is ANOTHER root of the tree tree.add_root(des_node) # a relation from the srouce node to the dest node v = RelationVisitor(src_id=relation_source_id, des_id=relation_dest_id, stix_objects=[], multi_root_tree=tree, log=logging) # visit the source node ret = source_node.accept(v) assert len(source_node.children) == 1 # dest node shall be moved as a child of source node assert source_node.children[0] == des_node # The tree should have only one root now assert len(tree.roots) == 1 # if we found, we shall stop assert not ret
def test(self): """ Use the following tree for test root1 child1 child2 child3 child4 child5 child6 child2(link) root2 child7 :return: """ root1 = TreeNode(obj_type="indicator", obj_id="", name="root1") root2 = TreeNode(obj_type="identity", obj_id="", name="root2") # # 1. First build a tree # tree = MultiRootTree(root1) tree.add_root(root2) child1 = TreeNode(obj_type="malware", obj_id="", name="child1") root1.add_child(child1) child2 = TreeNode(obj_type="indicator", obj_id="", name="child2") child1.add_child(child2) child3 = TreeNode(obj_type="malware", obj_id="", name="child3") child1.add_child(child3) child4 = TreeNode(obj_type="observed-data", obj_id="", objects={"0": { "type": "file", "name": "child4" }}, name="child4") child3.add_child(child4) child5 = TreeNode(obj_type="file", obj_id="", name="child5") root1.add_child(child5) child6 = TreeNode( obj_type="observed-data", obj_id="", objects={"0": { "type": "ipv4-addr", "value": "child6" }}, name="child6") child5.add_child(child6) child2_link = TreeNode(obj_type="indicator", obj_id="", name="child2") child2_link.is_link = True child5.add_child(child2_link) child7 = TreeNode(obj_type="ipv4-addr", obj_id="", name="child7") root2.add_child(child7) # # Use a html_gen_visitor to parse the tree # visitor = HtmlGenVisitor(logging) tree.accept(visitor=visitor) # # Output the html file # with open("TestGenHtml.html", "w") as outfile: outfile.write(visitor.html) # Visually verify that the output html # shows a tree similar to the one above assert "root1" in visitor.html assert "root2" in visitor.html for x in range(1, 7, 1): # assert each child is there name = "child" + str(x) assert name in visitor.html