Change command line arguments

myth

@@ -34,8 +34,7 @@ parser.add_argument('-d', '--disassemble',  action='store_true', help='disassemb
 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', 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('-g', '--graph', help='generate a call graph', 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')
@@ -54,7 +53,7 @@ contract_storage = get_persistent_storage(db_dir)
 
 args = parser.parse_args()
 
-if (args.disassemble):
+if (args.disassemble or args.graph):
 
     if (args.code):
         encoded_bytecode = args.code
@@ -69,7 +68,7 @@ if (args.disassemble):
             exitWithError("Exception loading bytecode via RPC: " + str(e))
 
     else:
-        exitWithError("Disassembler: Provide the input bytecode via -c BYTECODE or --id ID")
+        exitWithError("No input bytecode. Please provide the code via -c BYTECODE or -a address")
 
     try:
         disassembly = Disassembly(encoded_bytecode)
@@ -77,16 +76,20 @@ if (args.disassemble):
     except binascii.Error:
         exitWithError("Disassembler: Invalid code string.")
 
-    easm_text = disassembly.get_easm()
+    if (args.disassemble):
 
-    if (args.outfile):
-        util.string_to_file(args.outfile, easm_text)
-    else:
+        easm_text = disassembly.get_easm()
         sys.stdout.write(easm_text)
 
     if (args.graph):
-        generate_callgraph(disassembly, args.graph)
 
+        try:
+            disassembly = Disassembly(encoded_bytecode)
+            # instruction_list = asm.disassemble(util.safe_decode(encoded_bytecode))
+        except binascii.Error:
+            exitWithError("Disassembler: Invalid code string.")
+
+        generate_callgraph(disassembly, args.graph)
 
 elif (args.trace):
 

mythril/disassembler/callgraph.py

@@ -1,64 +1,9 @@
-import graphviz as gv
-
-
-styles = {
-    'graph': {
-        'overlap': 'false',
-        'fontsize': '16',
-        'fontcolor': 'white',
-        'bgcolor': '#333333',
-        'concentrate':'true',
-    },
-    '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
-
+from laser.ethereum import svm, cfg
 
 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)
+    _svm = svm.SVM(disassembly)
 
-	graph.render(file)
+    _svm.sym_exec()
 
+    cfg.generate_callgraph(_svm, file)

setup.py

@@ -219,7 +219,7 @@ security community.
 setup(
     name='mythril',
 
-    version='0.4',
+    version='0.',
 
     description='A reversing and bug hunting framework for the Ethereum blockchain',
     long_description=long_description,