From 5fe21f5ec39a7fe841bc59daa85c410336beae0b Mon Sep 17 00:00:00 2001
From: Bernhard Mueller <bernhard.mueller11@gmail.com>
Date: Fri, 19 Jan 2018 18:56:30 +0700
Subject: [PATCH] Add some error checking

---
 mythril/analysis/report.py |  5 ++-
 mythril/support/truffle.py | 70 +++++++++++++++++++++++++++++++++++---
 2 files changed, 69 insertions(+), 6 deletions(-)

diff --git a/mythril/analysis/report.py b/mythril/analysis/report.py
index eead5278..70fa96e6 100644
--- a/mythril/analysis/report.py
+++ b/mythril/analysis/report.py
@@ -11,6 +11,7 @@ class Issue:
         self.description = description
         self.type = _type
         self.debug = debug
+        self.code = None
 
 
     def as_dict(self):
@@ -47,10 +48,12 @@ class Report:
 
             text += issue.description + "\n--------------------\n"
 
+            if issue.code:
+                text += "Affected code:\n\n" + issue.code + "\n--------------------\n"
+
             if len(issue.debug):
                 text += "++++ Debugging info ++++\n" + issue.debug + "\n"
 
             text+="\n"
 
         return text
-
diff --git a/mythril/support/truffle.py b/mythril/support/truffle.py
index 734422ea..f1e276ae 100644
--- a/mythril/support/truffle.py
+++ b/mythril/support/truffle.py
@@ -1,5 +1,11 @@
 import os
+import re
 import json
+from mythril.ether import util
+from mythril.ether.ethcontract import ETHContract
+from mythril.analysis.security import fire_lasers
+from mythril.analysis.symbolic import StateSpace
+from laser.ethereum import helper
 
 
 def analyze_truffle_project():
@@ -13,11 +19,65 @@ def analyze_truffle_project():
     for contract_file in contract_files:
 
         with open(os.path.join(build_dir, contract_file)) as cf:
-            contract = json.load(cf)
+            contractdata = json.load(cf)
 
-            name = contract['contractName']
+            name = contractdata['contractName']
+            bytecode = contractdata['deployedBytecode']
+
+            if (len(bytecode) < 4):
+                continue
+
+            ethcontract= ETHContract(bytecode, name=name, address = util.get_indexed_address(0))
+
+            contracts = [ethcontract]
+
+            states = StateSpace(contracts, max_depth = 10)
+            report = fire_lasers(states)
+
+            # augment with source code
+
+            disassembly = ethcontract.get_disassembly()
+            source = contractdata['source']
+
+            deployedSourceMap = contractdata['deployedSourceMap'].split(";")
+
+            mappings = []
+            i = 0
+
+            while(i < len(deployedSourceMap)):
+
+                # print("INDEX: " + str(i))
+                # print(deployedSourceMap[i])
+
+                m = re.search(r"^(\d+):*(\d+)", deployedSourceMap[i])
+
+                if (m):
+                    offset = m.group(1)
+                    length = m.group(2)
+                else:
+                    m = re.search(r"^:(\d+)", deployedSourceMap[i])
+
+                    if m:
+                        length = m.group(1)
+
+                # print("APPEND: " + str((offset, length)))
+
+                mappings.append((int(offset), int(length)))
+
+                i += 1
+
+            for key, issue in report.issues.items():
+
+                index = helper.get_instruction_index(disassembly.instruction_list, issue.pc)
+
+                # print("INDEX: " + str(index))
+                # print(mappings[index])
+
+                if index:
+                    issue.code_start = mappings[index][0]
+                    issue.code_length = mappings[index][1]
+                    issue.code = source[mappings[index][0]: mappings[index][0] + mappings[index][1]]
+
+            print(report.as_text())
 
-            bytecode = contract['deployedBytecode']
-            deployedSourceMap = contract['deployedSourceMap']
 
-            
\ No newline at end of file