def test_add_node(self):
parent_node = TreeNode()
parent_node.log = logging
parent_node.add_child(None)
assert len(parent_node.children) == 0
def test_match_destination(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"
node = TreeNode()
# The node to visit matches the dest id
node.id = relation_dest_id
# This node is the root of the tree
tree = MultiRootTree(node)
# The dest node is in the objects list
src_type = "file"
src_name = "fake file name"
src_node = {
"id": relation_source_id,
"type": src_type,
"name": src_name
}
stix_objects = []
stix_objects.append(src_node)
v = RelationVisitor(src_id=relation_source_id,
des_id=relation_dest_id,
stix_objects=stix_objects,
multi_root_tree=tree,
log=logging)
ret = node.accept(v)
#
# A new tree node shall be inserted as the root
#
new_node = tree.roots[0]
assert len(new_node.children) == 1
assert new_node.id == relation_source_id
assert new_node.type == src_type
assert new_node.name == src_name
assert new_node.is_root
assert not new_node.is_link
# the dest node is the child
assert new_node.children[0] == node
# if we found, we shall stop
assert not ret
def test_not_match(self):
"""
If the node the visitor visits does not match either the source id or the destination
id, it should just return True to continue
:return:
"""
node = TreeNode()
node.id = "Node Id not match"
v = RelationVisitor("source id", "dest id", [], None, logging)
assert v.action(node)
def test_link_node(self):
"""
When visit a link node, the visitor.action should just return True to continue.
No relation will be added to a link node. The relation will be added to the
node it linked to.
:return:
"""
node = TreeNode()
node.is_link = True
v = RelationVisitor("", "", [], None, logging)
assert v.action(node)
def test_no_image_found(self):
"""
:return:
"""
visitor = HtmlGenVisitor(log=logging)
fake_name = "Fake name"
fake_type = "fake_type"
root1 = TreeNode(
obj_type="observed-data",
obj_id="",
objects={"0": {
"type": fake_type,
"value": fake_name
}},
name=fake_name)
tree = MultiRootTree(root1)
tree.accept(visitor)
#
# Even if we can't find the image associate with it, we shall still show the
# stix object
#
assert fake_name in visitor.html
assert fake_type in visitor.html
def test_extract_failure(self):
"""
Try to extract non-existing node
:return:
"""
root = TreeNode()
root.id = "root id"
tree = stix_tree.MultiRootTree(root)
obj = {"id": "Other obj id"}
stix_objects = [obj]
existing, node = stix_tree.extract_object(stix_objects, tree,
"id not exist", logging)
assert not existing
assert not node
def test_tree_node_copy(self):
orig_node = TreeNode(obj_type="fake type",
obj_id="fake id",
name="fake name",
description="fake desc",
objects={"0": {
"type": "fake obj type"
}},
is_link=True,
log=None)
copy_node = orig_node.copy()
assert orig_node.type == copy_node.type
assert orig_node.id == copy_node.id
assert orig_node.name == copy_node.name
assert orig_node.description == copy_node.description
assert orig_node.objects == copy_node.objects
def test_empty_observed_data(self):
visitor = HtmlGenVisitor(log=logging)
fake_name = "Fake name"
type = "observed-data"
root1 = TreeNode(obj_type=type, obj_id="", objects={}, name=fake_name)
tree = MultiRootTree(root1)
tree.accept(visitor)
#
# If the objects field of an observed-data is empyt, we show it as
# type=observed-data, and value=name
#
assert fake_name in visitor.html
assert type in visitor.html
def test_get_node_visitor_found(self):
obj_id = "object id to look for"
node_with_same_id = TreeNode(obj_type=None, obj_id=obj_id)
#
# this visitor looks for an object id different from the node
#
v = GetNodeVisitor(obj_id)
ret = v.action(node_with_same_id)
#
# if the id matches, we shall stop and store the node
#
assert ret == False
assert v.node == node_with_same_id
def test_get_node_visitor_not_found(self):
node_with_diff_id = TreeNode(obj_type=None, obj_id="fake id")
#
# this visitor looks for an object id different from the node
#
v = GetNodeVisitor("other obj id")
ret = v.action(node_with_diff_id)
#
# if the id does not match, we shall continue
#
assert ret == True
assert not v.node
def test_get_node_visitor_break(self):
"""
parent
child1 (with match id)
child2
:return:
"""
matched_obj_id = "matched obj id"
parent = TreeNode(obj_type=None, obj_id="fake id")
child1 = TreeNode(obj_type=None, obj_id=matched_obj_id)
child2 = TreeNode(obj_type=None, obj_id="fake id2")
parent.add_child(child1)
parent.add_child(child2)
v = GetNodeVisitor(matched_obj_id)
parent.accept(v)
assert v.node == child1
def test_no_value(self):
#
# This should not happen, but if a stix object has type, but not name,
# we show a type icon with empty value
#
visitor = HtmlGenVisitor(log=logging)
type = "file"
root1 = TreeNode(obj_type="observed-data",
obj_id="",
objects={"0": {
"type": type,
"value": ""
}},
name="")
tree = MultiRootTree(root1)
tree.accept(visitor)
assert type in visitor.html
def test_init_with_object(self):
"""
:return:
"""
objects = {"0": {"type": "file"}}
stix_obj = {
"id": "stix id",
"type": "domain-name",
"name": "mydomain.com",
"objects": objects
}
node = TreeNode().init_with_object(stix_obj)
assert node.type == stix_obj["type"]
assert node.name == stix_obj["name"]
assert node.objects == objects
assert node.is_link == False
assert node.is_root == False
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_source_link(self):
#
# A relation from relation_source_id to relation_dest_id
#
relation_source_id = "source object id"
relation_dest_id = "dest object id"
des_node = TreeNode()
des_node.id = relation_dest_id
src_node = TreeNode()
src_node.id = relation_source_id
# The tree has a parent node as the root and
# and both source node and dest node are children of the parent node
parent_node = TreeNode()
tree = MultiRootTree(parent_node)
parent_node.add_child(des_node)
parent_node.add_child(src_node)
v = RelationVisitor(src_id=relation_source_id,
des_id=relation_dest_id,
stix_objects=[],
multi_root_tree=tree,
log=logging)
#
# visit the src_node. Here both the src_node and the des_node
# are already in the tree
#
ret = src_node.accept(v)
#
# A new link node will be added to the src node
#
assert len(src_node.children) == 1
new_node = src_node.children[0]
# This new node is a copy of the des node
assert new_node.id == relation_dest_id
assert new_node.type == des_node.type
assert new_node.name == des_node.name
assert not new_node.is_root
assert new_node.is_link
# 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