mythril/mythril-repo-changes.diff

diff --git a/README.md b/README.md
index ab3d9fd..357a612 100644
--- a/README.md
+++ b/README.md
@@ -144,6 +144,14 @@ $ myth -ig ./graph.html -a 0x5c436ff914c458983414019195e0f4ecbef9e6dd --max-dept
 
 ~~The "bounce" effect, while awesome (and thus enabled by default), sometimes messes up the graph layout.~~ Try adding the `--enable-physics` flag for a very entertaining "bounce" effect that unfortunately completely destroys usability.
 
+## Statespace JSON for Traceview Explorer
+
+The `-j FILENAME` option dumps the statespace to json in the format that is required by the TraceView Explorer GUI
+
+```bash
+$ ./myth -ij ./statespace.json -a 0x5c436ff914c458983414019195e0f4ecbef9e6dd --max-depth 8
+```
+
 ## Blockchain exploration
 
 If you are planning to do batch operations or use the contract search features, running a [go-ethereum](https://github.com/ethereum/go-ethereum) node is recommended. Start your local node as follows:
diff --git a/myth b/myth
index b3289bc..08e44d9 100755
--- a/myth
+++ b/myth
@@ -29,6 +29,7 @@ from mythril.support.loader import DynLoader
 from mythril.exceptions import CompilerError, NoContractFoundError
 from mythril.analysis.symbolic import SymExecWrapper
 from mythril.analysis.callgraph import generate_graph
+from mythril.analysis.traceexplore import get_serializable_statespace
 from mythril.analysis.security import fire_lasers
 from mythril.analysis.report import Report
 
@@ -57,6 +58,7 @@ commands.add_argument('-g', '--graph', help='generate a control flow graph', met
 commands.add_argument('-x', '--fire-lasers', action='store_true', help='detect vulnerabilities, use with -c, -a or solidity file(s)')
 commands.add_argument('-t', '--truffle', action='store_true', help='analyze a truffle project (run from project dir)')
 commands.add_argument('-d', '--disassemble', action='store_true', help='print disassembly')
+commands.add_argument('-j', '--statespace-json', help='dumps the statespace json', metavar='OUTPUT_FILE')
 
 inputs = parser.add_argument_group('input arguments')
 inputs.add_argument('-c', '--code', help='hex-encoded bytecode string ("6060604052...")', metavar='BYTECODE')
@@ -129,7 +131,7 @@ with open(signatures_file) as f:
 
 # Parse cmdline args
 
-if not (args.search or args.init_db or args.hash or args.disassemble or args.graph or args.fire_lasers or args.storage or args.truffle):
+if not (args.search or args.init_db or args.hash or args.disassemble or args.graph or args.fire_lasers or args.storage or args.truffle or args.statespace_json):
     parser.print_help()
     sys.exit()
 
@@ -402,6 +404,21 @@ elif args.graph or args.fire_lasers:
             'markdown': report.as_markdown() or "The analysis was completed successfully. No issues were detected."
         }
         print(outputs[args.outform])
+        
+elif args.statespace_json: 
+    if not contracts:
+        exitWithError(args.outform, "input files do not contain any valid contracts")
+    
+    if args.dynld:
+        sym = SymExecWrapper(contracts[0], address, dynloader=DynLoader(eth), max_depth=args.max_depth)
+    else:
+        sym = SymExecWrapper(contracts[0], address, max_depth=args.max_depth)
+
+    try:
+        with open(args.statespace_json, "w") as f:
+            json.dump(get_serializable_statespace(sym), f)
+    except Exception as e:
+        exitWithError(args.outform, "Error saving json: " + str(e))
 
 else:
     parser.print_help()
diff --git a/mythril/analysis/traceexplore.py b/mythril/analysis/traceexplore.py
new file mode 100644
index 0000000..cbe0391
--- /dev/null
+++ b/mythril/analysis/traceexplore.py
@@ -0,0 +1,109 @@
+from z3 import Z3Exception, simplify
+from laser.ethereum.svm import NodeFlags
+import re
+
+colors = [
+    {'border': '#26996f', 'background': '#2f7e5b', 'highlight': {'border': '#fff', 'background': '#28a16f'}},
+    {'border': '#9e42b3', 'background': '#842899', 'highlight': {'border': '#fff', 'background': '#933da6'}},
+    {'border': '#b82323', 'background': '#991d1d', 'highlight': {'border': '#fff', 'background': '#a61f1f'}},
+    {'border': '#4753bf', 'background': '#3b46a1', 'highlight': {'border': '#fff', 'background': '#424db3'}},
+    {'border': '#26996f', 'background': '#2f7e5b', 'highlight': {'border': '#fff', 'background': '#28a16f'}},
+    {'border': '#9e42b3', 'background': '#842899', 'highlight': {'border': '#fff', 'background': '#933da6'}},
+    {'border': '#b82323', 'background': '#991d1d', 'highlight': {'border': '#fff', 'background': '#a61f1f'}},
+    {'border': '#4753bf', 'background': '#3b46a1', 'highlight': {'border': '#fff', 'background': '#424db3'}},
+]
+
+def get_serializable_statespace(statespace):
+    
+    nodes = []
+    edges = []
+    
+    color_map = {}
+    i = 0
+    for k in statespace.accounts:
+        color_map[statespace.accounts[k].contract_name] = colors[i]
+        i += 1
+    
+    for node_key in statespace.nodes:
+
+        node = statespace.nodes[node_key]
+
+        code = node.get_cfg_dict()['code']
+        code = re.sub("([0-9a-f]{8})[0-9a-f]+", lambda m: m.group(1) + "(...)", code)
+
+        if NodeFlags.FUNC_ENTRY in node.flags:
+            code = re.sub("JUMPDEST", node.function_name, code)
+
+        code_split = code.split("\\n")
+
+        truncated_code = code if (len(code_split) < 7) else "\\n".join(code_split[:6]) + "\\n(click to expand +)"
+
+        color = color_map[node.get_cfg_dict()['contract_name']]
+                
+        def get_state_accounts(state):
+            state_accounts = []
+            for key in state.accounts:
+                account = state.accounts[key].as_dict()
+                account.pop('code', None)
+                account['balance'] = str(account['balance'])
+                
+                storage = {}
+                for storage_key in account['storage']:
+                    storage[str(storage_key)] = str(account['storage'][storage_key])
+            
+                state_accounts.append({
+                    'address': key,
+                    'storage': storage
+                })  
+            return state_accounts          
+                
+        states = [{'machine': x.mstate.as_dict(), 'accounts': get_state_accounts(x)} for x in node.states]
+        
+        for state in states:
+            state['machine']['stack'] = [str(s) for s in state['machine']['stack']]
+            state['machine']['memory'] = [str(m) for m in state['machine']['memory']]
+        
+        truncated_code = truncated_code.replace('\\n', '\n')
+        code = code.replace('\\n', '\n')
+        
+        s_node = {
+            'id': str(node_key),
+            'func': str(node.function_name),
+            'label': truncated_code,
+            'code': code,
+            'truncated': truncated_code,
+            'states': states,
+            'color': color,
+            'instructions': code.split('\n')
+        }
+        
+        nodes.append(s_node)
+
+    for edge in statespace.edges:
+
+        if (edge.condition is None):
+            label = ""
+        else:
+
+            try:
+                label = str(simplify(edge.condition)).replace("\n", "")
+            except Z3Exception:
+                label = str(edge.condition).replace("\n", "")
+
+        label = re.sub("([^_])([\d]{2}\d+)", lambda m: m.group(1) + hex(int(m.group(2))), label)
+        code = re.sub("([0-9a-f]{8})[0-9a-f]+", lambda m: m.group(1) + "(...)", code)
+
+        s_edge = {
+            'from': str(edge.as_dict()['from']),
+            'to': str(edge.as_dict()['to']),
+            'arrows': 'to',
+            'label': label,
+            'smooth': { 'type': "cubicBezier" }
+        }
+        
+        edges.append(s_edge)
+
+    return {
+        'edges': edges,
+        'nodes': nodes
+    }