Merge pull request #386 from norhh/add_truffle_signatures

Add truffle signatures

mythril/mythril.py

@@ -180,7 +180,7 @@ class Mythril(object):
         config.set('defaults', 'dynamic_loading', 'infura')
 
     def analyze_truffle_project(self, *args, **kwargs):
-        return analyze_truffle_project(*args, **kwargs)  # just passthru for now
+        return analyze_truffle_project(self.sigs, *args, **kwargs)  # just passthru by passing signatures for now
 
     def _init_solc_binary(self, version):
         # Figure out solc binary and version

mythril/support/signatures.py

@@ -214,9 +214,17 @@ class SignatureDb(object):
         :param code: solidity source code
         :return: dictionary {sighash: function_signature}
         """
-        sigs = {}
 
         funcs = re.findall(r'function[\s]+(.*?\))', code, re.DOTALL)
+        return SignatureDb.get_sigs_from_functions(funcs)
+
+    @staticmethod
+    def get_sigs_from_functions(funcs):
+        """
+        :param funcs: accepts a list of functions
+        :return: their signature mappings
+        """
+        sigs = {}
         for f in funcs:
             f = re.sub(r'[\n]', '', f)
             m = re.search(r'^([A-Za-z0-9_]+)', f)
@@ -240,5 +248,5 @@ class SignatureDb(object):
                 signature = re.sub(r'\s', '', signature)
                 sigs["0x" + utils.sha3(signature)[:4].hex()] = signature
 
-        logging.debug("Signatures: parse soldiity found %d signatures" % len(sigs))
+        logging.debug("Signatures: found %d signatures after parsing" % len(sigs))
         return sigs

mythril/support/truffle.py

@@ -6,6 +6,7 @@ import json
 import logging
 from mythril.ether.ethcontract import ETHContract
 from mythril.ether.soliditycontract import SourceMapping
+from mythril.exceptions import CriticalError
 from mythril.analysis.security import fire_lasers
 from mythril.analysis.symbolic import SymExecWrapper
 from mythril.analysis.report import Report
@@ -14,7 +15,7 @@ from mythril.ether import util
 from mythril.laser.ethereum.util import get_instruction_index
 
 
-def analyze_truffle_project(args):
+def analyze_truffle_project(sigs, args):
 
     project_root = os.getcwd()
 
@@ -33,13 +34,17 @@ def analyze_truffle_project(args):
                 name = contractdata['contractName']
                 bytecode = contractdata['deployedBytecode']
                 filename = PurePath(contractdata['sourcePath']).name
-            except:
+                abi = contractdata['abi']
+            except KeyError:
                 print("Unable to parse contract data. Please use Truffle 4 to compile your project.")
                 sys.exit()
-
-            if (len(bytecode) < 4):
+            if len(bytecode) < 4:
                 continue
 
+            list_of_functions = parse_abi_for_functions(abi)
+            sigs.signatures.update(sigs.get_sigs_from_functions(list_of_functions))
+            sigs.write()
+
             ethcontract = ETHContract(bytecode, name=name)
 
             address = util.get_indexed_address(0)
@@ -47,7 +52,7 @@ def analyze_truffle_project(args):
             issues = fire_lasers(sym)
 
             if not len(issues):
-                if (args.outform == 'text' or args.outform == 'markdown'):
+                if args.outform == 'text' or args.outform == 'markdown':
                     print("# Analysis result for " + name + "\n\nNo issues found.")
                 else:
                     result = {'contract': name, 'result': {'success': True, 'error': None, 'issues': []}}
@@ -97,13 +102,23 @@ def analyze_truffle_project(args):
 
                     report.append_issue(issue)
 
-                if (args.outform == 'json'):
+                if args.outform == 'json':
 
                     result = {'contract': name, 'result': {'success': True, 'error': None, 'issues': list(map(lambda x: x.as_dict, issues))}}
                     print(json.dumps(result))
 
                 else:
-                    if (args.outform == 'text'):
+                    if args.outform == 'text':
                         print("# Analysis result for " + name + ":\n\n" + report.as_text())
-                    elif (args.outform == 'markdown'):
+                    elif args.outform == 'markdown':
                         print(report.as_markdown())
+
+
+def parse_abi_for_functions(abi):
+    funcs = []
+    for data in abi:
+        if data['type'] != 'function':
+            continue
+        args = '('+','.join([input['type'] for input in data['inputs']])+')'
+        funcs.append(data['name']+args)
+    return funcs