Fix the block edges in generated graphs

Author: zeroSteiner

crimson_forge/base.py

@@ -146,6 +146,9 @@ def size(self):
 		return len(self.bytes)
 
 class DiGraphBase(networkx.DiGraph):
+	def _graph_edges(self):
+		return self.edges
+
 	def _graphml_id(self, node):
 		return node
 
@@ -168,7 +171,7 @@ def _graphml_graph(self, parent, id_prefix=''):
 			if hasattr(node, 'to_digraph'):
 				digraph = node.to_digraph()
 				digraph._graphml_graph(element, id_prefix=xml_node_id + ':')
-		for parent_node, child_node in self.edges:
+		for parent_node, child_node in self._graph_edges():
 			element = ElementTree.SubElement(graph, 'edge', attrib={
 				'source': id_prefix + self._graphml_id(parent_node),
 				'target': id_prefix + self._graphml_id(child_node)

crimson_forge/segment.py

@@ -44,6 +44,7 @@
 import crimson_forge.tailor as tailor
 
 import angr
+import boltons.iterutils
 import capstone
 import keystone
 
@@ -78,12 +79,16 @@ def __init__(self, blocks, *args, **kwargs):
 		super(BlocksDiGraph, self).__init__(*args, **kwargs)
 		self._blocks = blocks
 		t_blocks = tuple(blk for blk in self._blocks if isinstance(blk, block.BasicBlock))
-
 		self.add_nodes_from(t_blocks)
 		for blk in t_blocks:
 			for child in blk.children.values():
 				self.add_edge(blk, child)
 
+	def _graph_edges(self):
+		# override this because block-edges are flow directions and not positional constraints like they are in other
+		# digraphs
+		return boltons.iterutils.pairwise(sorted(self.nodes, key=lambda block: block.address))
+
 	def _graphml_id(self, blk):
 		return "block.0x{:04x}".format(blk.address)