diff --git a/.gitignore b/.gitignore index f9dec49b..d080669a 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,4 @@ __pycache__ mythril.egg-info build dist +contracts.json diff --git a/contractstorage.py b/contractstorage.py new file mode 100644 index 00000000..e059335b --- /dev/null +++ b/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']) + diff --git a/ethcontract.py b/ethcontract.py new file mode 100644 index 00000000..4b27c6d4 --- /dev/null +++ b/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) + diff --git a/mythril b/mythril index beadb97c..024aba7b 100755 --- a/mythril +++ b/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()