From 647dda765e53f97eebec4ce061e826e94d231d6b Mon Sep 17 00:00:00 2001 From: Bernhard Mueller Date: Sat, 28 Oct 2017 16:57:43 +0700 Subject: [PATCH] Refactor --- myth | 94 +++++++++++-------------------- mythril/disassembler/callgraph.py | 4 +- mythril/ether/util.py | 4 ++ tests/svm_test.py | 7 ++- 4 files changed, 44 insertions(+), 65 deletions(-) diff --git a/myth b/myth index 16901d8d..6cde6ffe 100755 --- a/myth +++ b/myth @@ -77,6 +77,8 @@ try: except KeyError: solc_binary = 'solc' +# Parse cmdline args + args = parser.parse_args() if not (args.search or args.init_db or args.hash or args.disassemble or args.graph or args.xrefs or args.fire_lasers or args.trace): @@ -134,26 +136,26 @@ if (args.address or len(args.solidity_file)): except Exception as e: exitWithError("Error establishing RPC connection: " + str(e)) -# Load input contracts +# Load / compile input contracts contracts = [] if (args.code): - contracts.append(ETHContract(args.code, address = util.get_random_address())) + contracts.append(ETHContract(args.code, address = util.get_indexed_address(0))) elif (args.address): contracts.append(ETHContract(eth.eth_getCode(args.address), address = args.address)) elif (len(args.solidity_file)): - counter = 0 + index = 0 for file in args.solidity_file: file = file.replace("~", str(Path.home())) # Expand user path name, bytecode = compile_solidity(solc_binary, file) - # Max. 16 contracts supported + # Max. 16 contracts supported! - contract = ETHContract(bytecode, name = name, address = "0x" + hex(counter)[2:] * 40) - counter += 1 + contract = ETHContract(bytecode, name = name, address = util.get_indexed_address(index)) + index += 1 contracts.append(contract) logging.info(contract.name + " at " + contract.address) @@ -161,6 +163,7 @@ elif (len(args.solidity_file)): else: exitWithError("No input bytecode. Please provide EVM code via -c BYTECODE, -a ADDRESS, or -i SOLIDITY_FILES") + # Commands if (args.disassemble): @@ -168,79 +171,50 @@ if (args.disassemble): easm_text = contracts[0].get_easm() sys.stdout.write(easm_text) -elif (args.graph): - - if args.enable_physics is not None: - physics = True - - modules = [] - - for contract in contracts: - modules.append(contract.as_dict()) +elif (args.trace): - html = generate_callgraph(modules, args.enable_physics) + if (args.data): + trace = evm.trace(contracts[0].code, args.data) - try: - with open(args.graph, "w") as f: - f.write(html) - except Exception as e: + else: + trace = evm.trace(contracts[0].code) - print("Error saving graph: " + str(e)) + for i in trace: + if (re.match(r'^PUSH.*', i['op'])): + print(str(i['pc']) + " " + i['op'] + " " + i['pushvalue'] + ";\tSTACK: " + i['stack']) + else: + print(str(i['pc']) + " " + i['op'] + ";\tSTACK: " + i['stack']) elif (args.xrefs): - contract = ETHContract(bytecode) + print("\n".join(contracts[0].get_xrefs())) - print("\n".join(contract.get_xrefs())) +elif (args.graph) or (args.fire_lasers): -elif (args.fire_lasers): + # Convert to LASER SVM format - modules = [] + modules = {} for contract in contracts: - modules.append(contract.as_dict()) + modules[contract.address] = contract.as_dict() - laserfree.fire(modules) + if (args.graph): -elif (args.trace): - - if (args.code): - bytecode = args.code + if args.enable_physics is not None: + physics = True - elif (args.address): - if args.ipc: - eth = EthIpc() - bytecode = eth.eth_getCode(args.address) + html = generate_callgraph(modules, args.enable_physics) - else: - if args.infura_mainnet: - eth = EthJsonRpc('mainnet.infura.io', 443, True) - elif args.infura_rinkeby: - eth = EthJsonRpc('rinkeby.infura.io', 443, True) - elif args.infura_kovan: - eth = EthJsonRpc('kovan.infura.io', 443, True) - elif args.infura_ropsten: - eth = EthJsonRpc('ropsten.infura.io', 443, True) - else: - eth = EthJsonRpc(args.rpchost, args.rpcport, args.rpctls) - - bytecode = eth.eth_getCode(args.address) - - else: - exitWithError("Disassembler: Provide the input bytecode via -c BYTECODE or --id ID") + try: + with open(args.graph, "w") as f: + f.write(html) + except Exception as e: - if (args.data): - trace = evm.trace(bytecode, args.data) + print("Error saving graph: " + str(e)) else: - trace = evm.trace(bytecode) - - for i in trace: - if (re.match(r'^PUSH.*', i['op'])): - print(str(i['pc']) + " " + i['op'] + " " + i['pushvalue'] + ";\tSTACK: " + i['stack']) - else: - print(str(i['pc']) + " " + i['op'] + ";\tSTACK: " + i['stack']) + laserfree.fire(modules) else: parser.print_help() diff --git a/mythril/disassembler/callgraph.py b/mythril/disassembler/callgraph.py index ec05c91e..76589d73 100644 --- a/mythril/disassembler/callgraph.py +++ b/mythril/disassembler/callgraph.py @@ -128,9 +128,9 @@ def serialize(_svm): -def generate_callgraph(contracts, physics): +def generate_callgraph(modules, physics): - _svm = svm.SVM(contracts) + _svm = svm.SVM(modules) _svm.sym_exec() diff --git a/mythril/ether/util.py b/mythril/ether/util.py index d1add330..670ca33f 100644 --- a/mythril/ether/util.py +++ b/mythril/ether/util.py @@ -30,3 +30,7 @@ def encode_calldata(func_name, arg_types, args): def get_random_address(): return binascii.b2a_hex(os.urandom(20)).decode('UTF-8') + +def get_indexed_address(index): + return "0x" + (hex(index)[2:] * 40) + diff --git a/tests/svm_test.py b/tests/svm_test.py index 2b4fecde..e3c1e7d9 100644 --- a/tests/svm_test.py +++ b/tests/svm_test.py @@ -7,9 +7,10 @@ class SVMTestCase(unittest.TestCase): def runTest(self): - contract = {'name': 'coinCaller', 'address': '0x0000000000000000000000000000000000000000', 'creation_code': '', 'code': '606060405260043610610041576000357c0100000000000000000000000000000000000000000000000000000000900463ffffffff16806390b98a1114610046575b600080fd5b341561005157600080fd5b610086600480803573ffffffffffffffffffffffffffffffffffffffff16906020019091908035906020019091905050610088565b005b60008073cd912af5286dc0b0a23edb30209734ecb60838b391508173ffffffffffffffffffffffffffffffffffffffff1663412664ae85856000604051602001526040518363ffffffff167c0100000000000000000000000000000000000000000000000000000000028152600401808373ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff16815260200182815260200192505050602060405180830381600087803b151561014d57600080fd5b6102c65a03f1151561015e57600080fd5b505050604051805190509050505050505600a165627a7a723058207ee3ba7d05ac95dae526efd929324ab3185c3a50cb85e147e79bccc90db1eb360029'} - contract['disassembly'] = Disassembly(contract['code']) + modules = {} + modules['0x0000000000000000000000000000000000000000'] = {'name': 'metaCoin', 'address': '0x0000000000000000000000000000000000000000', 'creation_code': '', 'code': '60606040526004361061004c576000357c0100000000000000000000000000000000000000000000000000000000900463ffffffff16806327e235e314610051578063412664ae1461009e575b600080fd5b341561005c57600080fd5b610088600480803573ffffffffffffffffffffffffffffffffffffffff169060200190919050506100f8565b6040518082815260200191505060405180910390f35b34156100a957600080fd5b6100de600480803573ffffffffffffffffffffffffffffffffffffffff16906020019091908035906020019091905050610110565b604051808215151515815260200191505060405180910390f35b60006020528060005260406000206000915090505481565b6000816000803373ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff16815260200190815260200160002054101561016157600090506101fe565b816000803373ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff16815260200190815260200160002060008282540392505081905550816000808573ffffffffffffffffffffffffffffffffffffffff1673ffffffffffffffffffffffffffffffffffffffff16815260200190815260200160002060008282540192505081905550600090505b929150505600a165627a7a72305820fd4fa106da498514e90965a45ffecc1da53a0cd8bb7a7135910f8612245a46370029'} + modules['0x0000000000000000000000000000000000000000']['disassembly'] = Disassembly(modules['0x0000000000000000000000000000000000000000']['code']) - html = generate_callgraph([contract], False) + html = generate_callgraph(modules, False) self.assertTrue("var nodes = [\n{id: 0, size: 150" in html)