Add callgraphs

disassembler/callgraph.py

@@ -0,0 +1,63 @@
+import graphviz as gv
+
+
+styles = {
+    'graph': {
+        'overlap': 'false',
+        'fontsize': '16',
+        'fontcolor': 'white',
+        'bgcolor': '#333333',
+    },
+    'nodes': {
+        'fontname': 'Helvetica',
+        'shape': 'box',
+        'fontcolor': 'white',
+        'color': 'white',
+        'style': 'filled',
+        'fillcolor': '#006699',
+    },
+    'edges': {
+        'style': 'dashed',
+        'dir': 'forward',
+        'color': 'white',
+        'arrowhead': 'normal',
+        'fontname': 'Courier',
+        'fontsize': '12',
+        'fontcolor': 'white',
+    }
+}
+
+def apply_styles(graph, styles):
+    graph.graph_attr.update(
+        ('graph' in styles and styles['graph']) or {}
+    )
+    graph.node_attr.update(
+        ('nodes' in styles and styles['nodes']) or {}
+    )
+    graph.edge_attr.update(
+        ('edges' in styles and styles['edges']) or {}
+    )
+    return graph
+
+
+def generate_callgraph(disassembly, file):
+
+	graph = gv.Graph(format='svg')
+
+	index = 0
+
+	for block in disassembly.blocks:
+		easm = block.get_easm().replace("\n", "\l")
+
+		graph.node(str(index), easm)
+		index += 1
+
+	for xref in disassembly.xrefs:
+
+		graph.edge(str(xref[0]), str(xref[1]))
+
+
+	graph = apply_styles(graph, styles)
+
+	graph.render(file)
+

disassembler/disassembly.py

@@ -1,23 +1,28 @@
-from mythril.ether import asm,evm,util
-from mythril.rpc.client import EthJsonRpc
-from ethereum import utils
-import binascii
-import sys
+from mythril.ether import asm,util
 import os
 import json
+import sys
 
 
 class Block:
 
-    def __init__(self, code_index, start_addr, funcname):
+    def __init__(self, id, code_index, funcname):
+        self.id = id
         self.code_index = code_index
         self.funcname = funcname
-        self.xrefs = []
-
+        self.instruction_list = []
 
     def update_length(self, num_instructions):
         self.length = num_instructions
+        self.start_addr = self.instruction_list[0]['address']
+        self.end_addr = self.instruction_list[-1]['address']
 
+    def get_easm(self):
+
+        easm = str(self.instruction_list[0]['address']) + " " + self.funcname + "\n"
+        easm += asm.instruction_list_to_easm(self.instruction_list[1:])
+
+        return easm
 
 
 class Disassembly:
@@ -25,6 +30,7 @@ class Disassembly:
     def __init__(self, code):
         self.instruction_list = asm.disassemble(util.safe_decode(code))
         self.blocks = []
+        self.xrefs = []
         self.func_to_addr = {}
         self.addr_to_func = {}
 
@@ -60,25 +66,48 @@ class Disassembly:
 
         index = 0
         blocklen = 0
+        blocknumber = 1
 
         for instruction in self.instruction_list:
 
-            blocklen += 1
-
             if (instruction['opcode'] == "JUMPDEST"):
 
                 try:
-                    func_name = self.addr_to_func[instruction['address']]
+                    func_name = "- FUNCTION " + self.addr_to_func[instruction['address']] + " -"
                 except KeyError:
-                    func_name = "JUMPDEST_UNK"
+                    func_name = "- JUMPDEST_UNK -"
                 
                 current_block.update_length(blocklen)
                 self.blocks.append(current_block)
-                current_block = Block(index, instruction['address'], func_name)
+                current_block = Block(blocknumber, index, func_name)
                 blocklen = 0
+                blocknumber += 1
+                
+            current_block.instruction_list.append(instruction)
+            blocklen += 1
 
             index += 1
+        
+        # Resolve cross-references
+
+        for block in self.blocks:
+
+            jmp_indices = asm.find_opcode_sequence(["JUMP"], block.instruction_list)
+            jmp_indices += asm.find_opcode_sequence(["JUMPI"], block.instruction_list)
+
+            for i in jmp_indices:
+                try:
+                    dest_hex = block.instruction_list[i - 1]['argument']
+                    dest = int(dest_hex[2:], 16)
+                except:
+                    continue
+
+                j = 0
+
+                while(self.blocks[j].end_addr < dest):
+                    j += 1
 
+                self.xrefs.append((block.id, self.blocks[j].id))
 
 
     def get_easm(self):
@@ -86,8 +115,6 @@ class Disassembly:
         easm = asm.instruction_list_to_easm(self.instruction_list[0:self.blocks[0].length])
 
         for block in self.blocks[1:]:
-            easm += str(self.instruction_list[block.code_index]['address']) + " #### " + block.funcname + " ####\n"
-            
-            easm += asm.instruction_list_to_easm(self.instruction_list[block.code_index + 1:block.code_index + block.length])
+            easm += block.get_easm()
 
         return easm

myth

@@ -6,6 +6,7 @@
 
 from mythril.ether import evm,util
 from disassembler.disassembly import Disassembly
+from disassembler.callgraph import generate_callgraph
 from mythril.ether.contractstorage import get_persistent_storage
 from mythril.rpc.client import EthJsonRpc
 from ethereum import utils
@@ -30,11 +31,12 @@ def exitWithError(message):
 
 parser = argparse.ArgumentParser(description='Bug hunting on the Ethereum blockchain')
 
-parser.add_argument('-d', '--disassemble',  action='store_true', help='disassemble, use with -c or -a')
+parser.add_argument('-d', '--disassemble',  action='store_true', help='disassemble, specify input with -c or -a')
 parser.add_argument('-t', '--trace', action='store_true', help='trace, use with -c or -a and --data (optional)')
 parser.add_argument('-c', '--code', help='hex-encoded bytecode string ("6060604052...")', metavar='BYTECODE')
 parser.add_argument('-a', '--address', default='0x0123456789ABCDEF0123456789ABCDEF01234567', help='contract address')
 parser.add_argument('-o', '--outfile')
+parser.add_argument('-g', '--graph', help='when disassembling, also generate a callgraph', metavar='OUTPUT_FILE')
 parser.add_argument('--data', help='message call input data for tracing')
 parser.add_argument('--search', help='search the contract database')
 parser.add_argument('--xrefs', help='get xrefs from contract in database', metavar='CONTRACT_HASH')
@@ -66,14 +68,6 @@ if (args.disassemble):
 
         except Exception as e:
             exitWithError("Exception loading bytecode via RPC: " + str(e))
-    elif (args.infile):
-
-        try:
-
-            encoded_bytecode = util.file_to_string(args.infile).rstrip()
-
-        except Exception as e:
-            exitWithError("Exception loading bytecode from file: " + str(e))
 
     else:
         exitWithError("Disassembler: Provide the input bytecode via -c BYTECODE or --id ID")
@@ -91,6 +85,10 @@ if (args.disassemble):
     else:
         sys.stdout.write(easm_text)
 
+    if (args.graph):
+        generate_callgraph(disassembly, '/Users/berndt/Desktop/test.svg')
+
+
 elif (args.trace):
 
     if (args.code):

mythril/ether/asm.py

@@ -69,7 +69,7 @@ def find_opcode_sequence(pattern, instruction_list):
 
     pattern_length = len(pattern)
 
-    for i in range(0, len(instruction_list) - pattern_length):
+    for i in range(0, len(instruction_list) - pattern_length + 1):
 
         if instruction_list[i]['opcode'] == pattern[0]: