Graph: generic html dot code

4 files changed, 175 insertions, 117 deletions

diff --git a/example/disasm/full.py b/example/disasm/full.py
index 0b0069c6..ee0b88dd 100644
--- a/example/disasm/full.py
+++ b/example/disasm/full.py
@@ -148,7 +148,7 @@ for blocs in all_funcs_blocs.values():
 
 
 log.info('generate graph file')
-open('graph_execflow.dot', 'w').write(all_blocs.dot(label=True))
+open('graph_execflow.dot', 'w').write(all_blocs.dot(offset=True))
 
 log.info('generate intervals')
 
diff --git a/miasm2/core/asmbloc.py b/miasm2/core/asmbloc.py
index a4427206..4d7a197a 100644
--- a/miasm2/core/asmbloc.py
+++ b/miasm2/core/asmbloc.py
@@ -3,7 +3,6 @@
 
 import logging
 import inspect
-import re
 from collections import namedtuple
 
 import miasm2.expression.expression as m2_expr
@@ -254,6 +253,7 @@ class asm_bloc(object):
 
 
 class asm_block_bad(asm_bloc):
+
     """Stand for a *bad* ASM block (malformed, unreachable,
     not disassembled, ...)"""
 
@@ -261,7 +261,7 @@ class asm_block_bad(asm_bloc):
                    0: "Unable to disassemble",
                    1: "Null starting block",
                    2: "Address forbidden by dont_dis",
-    }
+                   }
 
     def __init__(self, label=None, alignment=1, errno=-1, *args, **kwargs):
         """Instanciate an asm_block_bad.
@@ -577,6 +577,7 @@ def dis_bloc_all(mnemo, pool_bin, offset, job_done, symbol_pool, dont_dis=[],
 
 
 class AsmCFG(DiGraph):
+
     """Directed graph standing for a ASM Control Flow Graph with:
      - nodes: asm_bloc
      - edges: constraints between blocks, synchronized with asm_bloc's "bto"
@@ -640,7 +641,7 @@ class AsmCFG(DiGraph):
     def add_uniq_edge(self, src, dst, constraint):
         """Add an edge from @src to @dst if it doesn't already exist"""
         if (src not in self._nodes_succ or
-            dst not in self._nodes_succ[src]):
+                dst not in self._nodes_succ[src]):
             self.add_edge(src, dst, constraint)
 
     def del_edge(self, src, dst):
@@ -704,76 +705,49 @@ class AsmCFG(DiGraph):
             # Use "_uniq_" beacause the edge can already exist due to add_node
             self.add_uniq_edge(*edge, constraint=graph.edges2constraint[edge])
 
-    def dot(self, label=False, lines=True):
-        """Render dot graph with HTML
-        @label: (optional) if set, add the corresponding label in each block
-        @lines: (optional) if set, includes assembly lines in the output
-        """
-
-        escape_chars = re.compile('[' + re.escape('{}') + ']')
-        label_attr = 'colspan="2" align="center" bgcolor="grey"'
-        edge_attr = 'label = "%s" color="%s" style="bold"'
-        td_attr = 'align="left"'
-        block_attr = 'shape="Mrecord" fontname="Courier New"'
-
-        out = ["digraph asm_graph {"]
-        fix_chars = lambda x: '\\' + x.group()
-
-        # Generate basic blocks
-        out_blocks = []
-        for block in self.nodes():
-            out_block = '%s [\n' % block.label.name
-            out_block += "%s " % block_attr
-            if isinstance(block, asm_block_bad):
-                out_block += 'style=filled fillcolor="red" '
-            out_block += 'label =<<table border="0" cellborder="0" cellpadding="3">'
-
-            block_label = '<tr><td %s>%s</td></tr>' % (
-                label_attr, block.label.name)
-            block_html_lines = []
-
-            if lines:
-                if isinstance(block, asm_block_bad):
-                    block_html_lines.append(block.ERROR_TYPES.get(block._errno,
-                                                                  block._errno))
-
-                for line in block.lines:
-                    if label:
-                        out_render = "%.8X</td><td %s> " % (line.offset,
-                                                            td_attr)
-                    else:
-                        out_render = ""
-                    out_render += escape_chars.sub(fix_chars, str(line))
-                    block_html_lines.append(out_render)
-
-            block_html_lines = ('<tr><td %s>' % td_attr +
-                                ('</td></tr><tr><td %s>' % td_attr).join(block_html_lines) +
-                                '</td></tr>')
-            out_block += "%s " % block_label
-            out_block += block_html_lines + "</table>> ];"
-            out_blocks.append(out_block)
+    def node2lines(self, node):
+        yield self.DotCellDescription(text=str(node.label.name),
+                                      attr={'align': 'center',
+                                      'colspan': 2,
+                                            'bgcolor': 'grey'})
+
+        if isinstance(node, asm_block_bad):
+            yield [self.DotCellDescription(
+                text=node.ERROR_TYPES.get(node._errno,
+                                          node._errno),
+                attr={})
+            raise StopIteration
+        for line in node.lines:
+            if self._dot_offset:
+                yield [self.DotCellDescription(text="%.8X" % line.offset,
+                                               attr={}),
+                       self.DotCellDescription(text=str(line), attr={})]
+            else:
+                yield self.DotCellDescription(text=str(line), attr={})
 
-        out += out_blocks
+    def node_attr(self, node):
+        if isinstance(node, asm_block_bad):
+            return {'style': 'filled', 'fillcolor': 'red'}
+        return {}
 
-        # Generate links
-        for src, dst in self.edges():
-            exp_label = dst.label
-            cst = self.edges2constraint.get((src, dst), None)
+    def edge_attr(self, src, dst):
+        cst = self.edges2constraint.get((src, dst), None)
+        edge_color = "blue"
 
-            edge_color = "black"
+        if len(self.successors(src)) > 1:
             if cst == asm_constraint.c_next:
                 edge_color = "red"
-            elif cst == asm_constraint.c_to:
+            else:
                 edge_color = "limegreen"
-            # special case
-            if len(src.bto) == 1:
-                edge_color = "blue"
 
-            out.append('%s -> %s' % (src.label.name, dst.label.name) + \
-                       '[' + edge_attr % (cst, edge_color) + '];')
+        return {"color": edge_color}
 
-        out.append("}")
-        return '\n'.join(out)
+    def dot(self, offset=False):
+        """
+        @offset: (optional) if set, add the corresponding offsets in each node
+        """
+        self._dot_offset = offset
+        return super(AsmCFG, self).dot()
 
     # Helpers
     @property
@@ -935,7 +909,8 @@ class AsmCFG(DiGraph):
         new block destinations
         @kwargs: (optional) named arguments to pass to dis_block_callback
         """
-        # Get all possible destinations not yet resolved, with a resolved offset
+        # Get all possible destinations not yet resolved, with a resolved
+        # offset
         block_dst = [label.offset
                      for label in self.pendings
                      if label.offset is not None]
diff --git a/miasm2/core/graph.py b/miasm2/core/graph.py
index 310adc2e..f49e0da1 100644
--- a/miasm2/core/graph.py
+++ b/miasm2/core/graph.py
@@ -1,9 +1,14 @@
 from collections import defaultdict, namedtuple
+import re
 
 
 class DiGraph(object):
     """Implementation of directed graph"""
 
+    # Stand for a cell in a dot node rendering
+    DotCellDescription = namedtuple("DotCellDescription",
+                                    ["text", "attr"])
+
     def __init__(self):
         self._nodes = set()
         self._edges = []
@@ -144,35 +149,100 @@ class DiGraph(object):
                     out.append(path + [dst])
         return out
 
+    def nodeid(self, node):
+        """
+        Returns uniq id for a @node
+        @node: a node of the graph
+        """
+        return hash(node) & 0xFFFFFFFFFFFFFFFF
+
+    def node2lines(self, node):
+        """
+        Returns an iterator on cells of the dot @node.
+        A DotCellDescription or a list of DotCellDescription are accepted
+        @node: a node of the graph
+        """
+        yield self.DotCellDescription(text=str(node), attr={})
+
+    def node_attr(self, node):
+        """
+        Returns a dictionnary of the @node's attributes
+        @node: a node of the graph
+        """
+        return {}
+
+    def edge_attr(self, src, dst):
+        """
+        Return a dictionnary of attributes for the edge between @src and @dst
+        @src: the source node of the edge
+        @dst: the destination node of the edge
+        """
+        return {}
+
     @staticmethod
-    def node2str(node):
-        return str(node)
+    def _fix_chars(token):
+        return "&#%04d;" % ord(token.group())
 
     @staticmethod
-    def edge2str(src, dst):
-        return ""
+    def _attr2str(default_attr, attr):
+        return ' '.join('%s="%s"' % (name, value)
+                        for name, value in
+                        dict(default_attr,
+                             **attr).iteritems())
 
     def dot(self):
-        out = """
-digraph asm_graph {
-graph [
-splines=polyline,
-];
-node [
-fontsize = "16",
-shape = "box"
-];
-"""
+        """Render dot graph with HTML"""
+
+        escape_chars = re.compile('[' + re.escape('{}') + '&|<>' + ']')
+        label_attr = 'colspan="2" align="center" bgcolor="grey"'
+        edge_attr = 'label = "%s" color="%s" style="bold"'
+        td_attr = {'align': 'left'}
+        nodes_attr = {'shape': 'Mrecord',
+                      'fontname': 'Courier New'}
+
+        out = ["digraph asm_graph {"]
+
+        # Generate basic nodes
+        out_nodes = []
         for node in self.nodes():
-            out += '%s [label="%s"];\n' % (
-                hash(node) & 0xFFFFFFFFFFFFFFFF, self.node2str(node))
+            node_id = self.nodeid(node)
+            out_node = '%s [\n' % node_id
+            out_node += self._attr2str(nodes_attr, self.node_attr(node))
+            out_node += 'label =<<table border="0" cellborder="0" cellpadding="3">'
+
+            node_html_lines = []
+
+            for lineDesc in self.node2lines(node):
+                out_render = ""
+                if isinstance(lineDesc, self.DotCellDescription):
+                    lineDesc = [lineDesc]
+                for col in lineDesc:
+                    out_render += "<td %s>%s</td>" % (
+                        self._attr2str(td_attr, col.attr),
+                        escape_chars.sub(self._fix_chars, str(col.text)))
+                node_html_lines.append(out_render)
+
+            node_html_lines = ('<tr>' +
+                               ('</tr><tr>').join(node_html_lines) +
+                               '</tr>')
+
+            out_node += node_html_lines + "</table>> ];"
+            out_nodes.append(out_node)
 
+        out += out_nodes
+
+        # Generate links
         for src, dst in self.edges():
-            out += '%s -> %s [label="%s"]\n' % (hash(src) & 0xFFFFFFFFFFFFFFFF,
-                                                hash(dst) & 0xFFFFFFFFFFFFFFFF,
-                                                self.edge2str(src, dst))
-        out += "}"
-        return out
+            attrs = self.edge_attr(src, dst)
+
+            attrs = ' '.join('%s="%s"' % (name, value)
+                             for name, value in attrs.iteritems())
+
+            out.append('%s -> %s' % (self.nodeid(src), self.nodeid(dst)) +
+                       '[' + attrs + '];')
+
+        out.append("}")
+        return '\n'.join(out)
 
     @staticmethod
     def _reachable_nodes(head, next_cb):
diff --git a/miasm2/ir/ir.py b/miasm2/ir/ir.py
index 9d90f79a..35054f50 100644
--- a/miasm2/ir/ir.py
+++ b/miasm2/ir/ir.py
@@ -130,36 +130,49 @@ class DiGraphIR(DiGraph):
         self._blocks = blocks
         super(DiGraphIR, self).__init__(*args, **kwargs)
 
-    def dot(self):
-        """Output the graphviz script"""
-        out = """
-    digraph asm_graph {
-    size="80,50";
-    node [
-    fontsize = "16",
-    shape = "box"
-    ];
+    def node2lines(self, node):
+        yield self.DotCellDescription(text=str(node.name),
+                                      attr={'align': 'center',
+                                      'colspan': 2,
+                                            'bgcolor': 'grey'})
+        if node not in self._blocks:
+            yield [self.DotCellDescription(text="NOT PRESENT", attr={})]
+            raise StopIteration
+        for i, exprs in enumerate(self._blocks[node].irs):
+            for expr in exprs:
+                if self._dot_offset:
+                    yield [self.DotCellDescription(text="%-4d" % i, attr={}),
+                           self.DotCellDescription(text=str(expr), attr={})]
+                else:
+                    yield self.DotCellDescription(text=str(expr), attr={})
+            yield self.DotCellDescription(text="", attr={})
+
+    def edge_attr(self, src, dst):
+        if src not in self._blocks or dst not in self._blocks:
+            return {}
+        src_irdst = self._blocks[src].dst
+        edge_color = "blue"
+        if isinstance(src_irdst, m2_expr.ExprCond):
+            if (expr_is_label(src_irdst.src1) and
+                    src_irdst.src1.name == dst):
+                edge_color = "limegreen"
+            elif (expr_is_label(src_irdst.src2) and
+                  src_irdst.src2.name == dst):
+                edge_color = "red"
+        return {"color": edge_color}
+
+    def node_attr(self, node):
+        if node not in self._blocks:
+            return {'style': 'filled', 'fillcolor': 'red'}
+        return {}
+
+    def dot(self, offset=False):
         """
-        all_lbls = {}
-        for lbl in self.nodes():
-            if lbl not in self._blocks:
-                continue
-            irb = self._blocks[lbl]
-            ir_txt = [str(lbl)]
-            for irs in irb.irs:
-                for l in irs:
-                    ir_txt.append(str(l))
-                ir_txt.append("")
-            ir_txt.append("")
-            all_lbls[hash(lbl)] = "\l\\\n".join(ir_txt)
-        for l, v in all_lbls.items():
-            out += '%s [label="%s"];\n' % (l, v)
-
-        for a, b in self.edges():
-            out += '%s -> %s;\n' % (hash(a), hash(b))
-        out += '}'
-        return out
-
+        @offset: (optional) if set, add the corresponding line number in each
+        node
+        """
+        self._dot_offset = offset
+        return super(DiGraphIR, self).dot()
 
 class ir(object):