def copy_zss_tree_with_replacement(t: zss.Node,
replace_map: dict,
max_depth=10):
# deepcopy is very slow, it becomes bottleneck
children = list()
if max_depth > 0 and len(t.children) > 0:
for c in t.children:
children.append(
copy_zss_tree_with_replacement(c, replace_map, max_depth - 1))
if t.label in replace_map.keys():
root = zss.Node(label=replace_map[t.label], children=children)
else:
root = zss.Node(label=t.label, children=children)
return root
def tree_to_zss(self, t):
if isinstance(t, int):
return zss.Node(str(t))
elif isinstance(t, list):
node = zss.Node('**LIST**')
for c in t:
node.addkid(self.tree_to_zss(c))
return node
elif isinstance(t, dict):
node = zss.Node('**DICT**')
for k in sorted(t.keys()):
child = zss.Node(k, [self.tree_to_zss(t[k])])
node.addkid(child)
return node
elif t is None:
return zss.Node('**NONE**')
assert False, "don't know what to do with %s" % t
def convert_obj_to_node(obj):
obj_label = get_obj_label(obj)
obj_node = zss.Node(obj_label)
params = list(obj.get_params().items())
# sort by param name to guarantee order
params = sorted(params, key=lambda x: x[0])
for param_name, param_obj in params:
if isinstance(param_obj, sklearn.pipeline.Pipeline):
param_node = to_tree(param_obj)
elif sklearn.base.is_classifier(param_obj):
param_node = convert_obj_to_node(param_obj)
else:
param_node = zss.Node("{}_{}".format(param_name, param_obj))
obj_node.addkid(param_node)
return obj_node
def to_tree(pipeline):
root = zss.Node("root")
curr = root
if isinstance(pipeline, sklearn.pipeline.FeatureUnion):
steps = pipeline.transformer_list
else:
steps = pipeline.steps
for step_name, step_obj in steps:
step_node = convert_obj_to_node(step_obj)
curr.addkid(step_node)
curr = step_node
return root
def copy_zss_tree_with_replacement_no_recursion(t: zss.Node,
replace_map: dict):
if t.label in replace_map.keys():
root = zss.Node(label=replace_map[t.label])
else:
root = zss.Node(label=t.label)
stack = [(c, 1) for c in reversed(t.children)]
cur_path = [root]
while len(stack) > 0:
tmp, depth = stack.pop()
if len(tmp.children) > 0:
cur_depth = len(cur_path)
for c in reversed(tmp.children):
stack.append((c, cur_depth))
cur_path = cur_path[:depth]
if tmp.label in replace_map.keys():
new_node = zss.Node(label=replace_map[tmp.label])
else:
new_node = zss.Node(label=tmp.label)
cur_path[-1].addkid(new_node)
cur_path.append(new_node)
return root
def _kernel_tree_edit_distance_create_tree(self, current_node):
"""
Convert the state to tree format required by zss library for calculating
tree edit distance
"""
def get_label(s, n):
"""
Returns the sort string and label for a tree node
label is symbol in the node:voxel coordinates
sort string also containt the rule number used at that node.
sort string is used to sort the nodes of a tree into a canonical order.
"""
voxel = s.spatial_model.voxels.get(n)
# if there are no voxels associated with a node, or it is a root node
if voxel is None or s.tree.nodes[n].bpointer is None:
voxel = [0, 0, 0]
voxel[1] = 0
sort_str = s.tree.nodes[n].tag.symbol + str(
s.tree.nodes[n].tag.rule)
label = s.tree.nodes[n].tag.symbol
voxel = ''.join([str(i) for i in voxel])
sort_str = sort_str + ':' + voxel
label = label + ':' + voxel
return sort_str, label
# get the sort strings for the children of this node
child_sort_strs = [
get_label(self, child)[0]
for child in self.tree.nodes[current_node].fpointer
]
# get the permutation we need to get children into sorted order
sort_indices = [
i[0]
for i in sorted(enumerate(child_sort_strs), key=lambda x: x[1])
]
# apply that permutation
children = [
self.tree.nodes[current_node].fpointer[i] for i in sort_indices
]
# call create_tree recursively on the children
t = zss.Node(get_label(self, current_node)[1],
children=[
self._kernel_tree_edit_distance_create_tree(cn)
for cn in children
])
return t
def build_tree(terms):
root = zss.Node("root")
node_list = [root]
left = 0
for term in terms:
if term == '(':
left += 1
elif term == ')':
left -= 1
node_list.pop()
else:
node = build_node(term)
father = node_list[-1]
father.addkid(node)
node_list.append(node)
assert left == 0 and len(root.children) == 1
return root.children[0]
def make_tree(list):
root = zss.Node('root')
kid1 = zss.Node('tmp')
kid2 = zss.Node('tmp')
kid3 = zss.Node('tmp')
for x in list:
l = len(x.split('.'))
if(l==1):
kid1 = zss.Node('1')
root.addkid(kid1)
elif(l==2):
kid2 = zss.Node('2')
kid1.addkid(kid2)
elif(l==3):
kid3 = zss.Node('3')
kid2.addkid(kid3)
return root
def zssTreeFromTreeLibTree(tree):
zss_root = zss.Node(tree.get_node(tree.root))
appendChildren(tree.get_node(tree.root), zss_root)
return zss_root
def appendChildren(treelib_node, zss_node):
# Get the children of the node
for child in treelib_node.get_children():
zss_child_node = zss.Node(child.tag)
appendChildren(child, zss_child_node)
zss_node.addkid(zss_child_node)
def convert_structure_to_zss(tree):
lbl = label(tree[0]) if hasattr(tree[0], "ftag") else str(type(tree[0]))
return zss.Node(lbl, [convert_structure_to_zss(n) for n in tree[1]])
def builupTree(root):
tree = zss.Node(root.nodeName)
for child in root.children:
tree.addkid(builupTree(child))
return tree
def builupTree(root, useNCD):
tree = zss.Node(root.info if useNCD else root.nodeName)
for child in root.children:
tree.addkid(builupTree(child, useNCD))
return tree
# Ground Truth Tree construction:
gt_graph, artificial_edge_list = createGraphFromTruth(ground_truth)
if save_plots:
# Save the Plot:
gt_file_name = ground_truth.split("/")[-1].split(".")[0]
plt.figure(figsize=(300, 300))
nx.draw(gt_graph, with_labels=True)
plt.savefig(plot_dir + gt_file_name + ".png")
gt_tree = nx.bfs_tree(gt_graph, source="root")
gt_nodes_dict = {}
for edge in gt_tree.edges():
if edge[0] not in gt_nodes_dict:
gt_nodes_dict[edge[0]] = zss.Node(edge[0])
if edge[1] not in gt_nodes_dict:
gt_nodes_dict[edge[1]] = zss.Node(edge[1])
gt_nodes_dict[edge[0]].addkid(gt_nodes_dict[edge[1]])
# Generated Tree construction:
gen_graph = nx.read_weighted_edgelist(mst_file)
gen_graph.add_edges_from(artificial_edge_list)
if save_plots:
# Save the Plot:
plt.figure(figsize=(300, 300))
nx.draw(gen_graph, with_labels=True)
gen_file_name = mst_file.split("/")[-1].split(".")[0]
plt.savefig(plot_dir + gen_file_name + ".png")
def build_node(info):
return zss.Node(info)
def expr2tree(e, exe_ranges, data_ranges):
if isinstance(e, int) or isinstance(e, float):
return zss.Node(label=e)
elif isinstance(e, claripy.bv.BVV):
value = e.value
if in_ranges(value, exe_ranges):
return zss.Node(label='exe_ptr')
elif in_ranges(value, data_ranges):
return zss.Node(label='data_ptr')
else:
return zss.Node(label=value)
elif isinstance(e, claripy.ast.base.Base):
if e.depth == 1:
if isinstance(e, claripy.Bits):
if e.symbolic:
# return zss.Node(label='vs' + str(e.size()))
return zss.Node(label=str(e))
elif isinstance(e, claripy.ast.bv.BV):
# here will receive all constants
# we separate all constants value to 3 types:
# executable_pointer, data_pointer, and normal constant
value = e.args[0]
if in_ranges(value, exe_ranges):
return zss.Node(label='exe_ptr')
elif in_ranges(value, data_ranges):
return zss.Node(label='data_ptr')
else:
return zss.Node(label=e.args[0])
else:
return zss.Node(label=str(e))
elif isinstance(e, claripy.ast.bool.Bool):
if e.symbolic:
return zss.Node(label='bs')
else:
return zss.Node(label='b')
else:
assert False, 'unknown base type %s' % e.__class__
elif e.depth == 2 and e.op == 'Extract':
# this is also a leaf of AST
return zss.Node(label=str(e.args[2]))
elif e.op == 'Concat':
# we ignore all concat 0 at head
tmp1 = expr2tree(e.args[0], exe_ranges, data_ranges)
if tmp1.label == 0:
return expr2tree(e.args[1], exe_ranges, data_ranges)
else:
node = zss.Node(label=e.op)
node.addkid(tmp1)
node.addkid(expr2tree(e.args[1], exe_ranges, data_ranges))
return node
else:
# if isinstance(e, claripy.ast.bv.BV) and e.depth == 2 and e.symbolic and e.op == 'Extract':
# return zss.Node(label='vs' + str(e.size()))
# else:
node = zss.Node(label=e.op)
for arg in e.args:
node.addkid(expr2tree(arg, exe_ranges, data_ranges))
return node
else:
assert False, 'unknown type %s' % e.__class__