Add contract search

.gitignore

@@ -7,3 +7,4 @@ __pycache__
 mythril.egg-info
 build
 dist
+contracts.json

contractstorage.py

@@ -0,0 +1,81 @@
+from rpc.client import EthJsonRpc
+from ethcontract import ETHContract
+from tinydb import TinyDB, Query
+import codecs
+import hashlib
+
+
+class ContractStorage:
+
+    def __init__(self):
+        self.db = TinyDB('./contracts.json')
+
+
+    def initialize(self, rpchost, rpcport):
+
+        eth = EthJsonRpc(rpchost, rpcport)
+
+        blockNum = eth.eth_blockNumber()
+
+        while(blockNum > 0):
+
+            if not blockNum % 1000:
+                print("Processing block: " + str(blockNum))
+
+            block = eth.eth_getBlockByNumber(blockNum)
+
+            for tx in block['transactions']:
+
+                if not tx['to']:
+
+                    receipt = eth.eth_getTransactionReceipt(tx['hash'])
+
+                    contract_address = receipt['contractAddress']
+
+                    contract_code = eth.eth_getCode(contract_address)
+
+                    m = hashlib.md5()
+
+                    m.update(contract_code.encode('UTF-8'))
+
+                    contract_hash = codecs.encode(m.digest(), 'hex_codec')
+                    contract_id = contract_hash.decode("utf-8")
+
+                    contract_balance = eth.eth_getBalance(contract_address)
+
+                    Contract = Query()
+
+                    new_instance = {'address': contract_address, 'balance': contract_balance}
+
+                    s = self.db.search(Contract.id == contract_id)
+
+                    if not len(s):
+
+                        self.db.insert({'id': contract_id, 'code': contract_code, 'instances': [new_instance]})
+
+                    else:
+
+                        instances = s[0]['instances']
+
+                        instances.append(new_instance)
+
+                        self.db.update({'instances': instances}, Contract.id == contract_id)
+
+            blockNum -= 1
+
+
+
+    def search(self, expression):
+
+        all_contracts = self.db.all()
+
+        for c in all_contracts:
+
+            for instance in c['instances']:
+
+                if ('balance' in instance):
+                    contract = ETHContract(c['code'], instance['balance'])
+
+                    if (contract.matches_expression(expression)):
+                        print("Found contract:" + instance['address'])
+

ethcontract.py

@@ -0,0 +1,38 @@
+from ether import asm, util
+import re
+
+
+
+class ETHContract:
+
+    def __init__(self, code = "", balance = 0):
+
+        self.easm_code = asm.disassembly_to_easm(asm.disassemble(util.safe_decode(code)))
+        self.balance = balance
+
+    def matches_expression(self, expression):
+
+        str_eval = ""
+
+        tokens = re.split("( and | or )", expression, re.IGNORECASE)
+
+        for token in tokens:
+
+            if token == " and " or token == " or ":
+                str_eval += token
+                continue
+
+            m = re.match(r'^code\(([a-zA-Z0-9\s,]+)\)$', token)
+
+            if (m):
+                code = m.group(1).replace(",", "\\n")
+                str_eval += "\"" + code + "\" in self.easm_code"
+                continue
+
+            m = re.match(r'^balance\s*[=><]+\s*\d+$', token)
+            if (m):
+                str_eval += "self." + m.group(0)
+                continue
+
+        return eval(str_eval)
+

mythril

@@ -5,7 +5,7 @@
 """
 
 from ether import asm,evm,util
-from rpc.client import EthJsonRpc
+from contractstorage import ContractStorage
 import sys
 import codecs
 import argparse
@@ -25,12 +25,15 @@ parser.add_argument('-c', '--code', help='hex-encoded bytecode string ("60606040
 parser.add_argument('-o', '--outfile')
 parser.add_argument('-f', '--infile', metavar='INPUTFILE')
 parser.add_argument('--txid', help='id of contract creation transaction')
+parser.add_argument('--init-db', action='store_true', help='Initialize the contract database')
+parser.add_argument('--search', help='search contract database')
 parser.add_argument('--rpchost', default='127.0.0.1', help='RPC host')
 parser.add_argument('--rpcport', type=int, default=8545, help='RPC port')
-parser.add_argument('--find-easm', help='search blockchain for opcode sequence')
 
 args = parser.parse_args()
 
+storage = ContractStorage()
+
 
 if (args.disassemble):
 
@@ -99,42 +102,12 @@ elif (args.trace):
 
     evm.trace(bytecode)
 
-elif (args.find_easm):
-
-    eth = EthJsonRpc(args.rpchost, args.rpcport)
-
-    searchstring = args.find_easm.replace(',',"\n")
-
-    blockNum = eth.eth_blockNumber()
-
-    while(blockNum > 0):
-
-        print(blockNum)
-
-        block = eth.eth_getBlockByNumber(blockNum)
-
-        for tx in block['transactions']:
-
-            if not tx['to']:
-
-                try:
-                    encoded_bytecode = util.bytecode_from_blockchain(tx['hash'], args.rpchost, args.rpcport)
-
-                    disassembly = asm.disassemble(util.safe_decode(encoded_bytecode))
-
-                    easm_text = asm.disassembly_to_easm(disassembly)
-
-                except RuntimeError as e:
-                    # Error getting contract code from transaction
-                    pass
-
-                if searchstring in easm_text:
-
-                    receipt = eth.eth_getTransactionReceipt(tx['hash'])
+elif (args.search):
 
-                    print("Found contract " + receipt['contractAddress'] + " created by transaction " + tx['hash'])
+    storage.search(args.search)
 
-        blockNum -= 1
+elif (args.init_db):
+        storage.initialize(args.rpchost, args.rpcport)  
 
 else:
     parser.print_help()