From cf62c49ff5d27de4da96becb3856912532f7315b Mon Sep 17 00:00:00 2001 From: Bernhard Mueller Date: Fri, 27 Oct 2017 15:23:21 +0700 Subject: [PATCH] Load contracts into list --- myth | 63 ++++++++++++++---------------------- mythril/ether/ethcontract.py | 6 ++-- mythril/ether/util.py | 11 +++++-- tests/ethcontract_test.py | 2 +- 4 files changed, 37 insertions(+), 45 deletions(-) diff --git a/myth b/myth index d66d8f54..8741e1a8 100755 --- a/myth +++ b/myth @@ -4,7 +4,7 @@ http://www.github.com/b-mueller/mythril """ -from mythril.ether import evm +from mythril.ether import evm, util from mythril.disassembler.disassembly import Disassembly from mythril.disassembler.callgraph import generate_callgraph from mythril.ether.contractstorage import get_persistent_storage @@ -48,8 +48,7 @@ commands.add_argument('--init-db', action='store_true', help='initialize the con inputs = parser.add_argument_group('input arguments') inputs.add_argument('-c', '--code', help='hex-encoded bytecode string ("6060604052...")', metavar='BYTECODE') -inputs.add_argument('-i', '--infile', help='solidity input file', metavar='INPUT_FILE') -inputs.add_argument('-l', '--libs', help='comma-separated list of solidity library files', metavar='LIBRARY_FILES') +inputs.add_argument('-i', '--infiles', help='comma-separated list of solidity files', metavar='INPUT_FILES') inputs.add_argument('-a', '--address', help='pull contract from the mainnet', metavar='CONTRACT_ADDRESS') inputs.add_argument('--data', help='message call input data for tracing') @@ -108,7 +107,7 @@ if args.search or args.init_db: # Establish RPC/IPC connection if necessary -if (args.address or args.infile): +if (args.address or args.infiles): if args.ipc: try: eth = EthIpc() @@ -122,48 +121,36 @@ if (args.address or args.infile): except Exception as e: exitWithError("Error establishing RPC connection: " + str(e)) -# Handle libs +# Load code -libraries = {} +contracts = [] -if (args.libs): - files = args.libs.split(",") +if (args.code): + contracts.append(ETHContract(args.code)) +elif (args.address): + contracts.append(ETHContract(eth.eth_getCode(args.address))) - c = 0 +elif (args.infiles): + files = args.infiles.split(",") for file in files: - address = binascii.b2a_hex(os.urandom(20)).decode('UTF-8') - bytecode = compile_solidity(solc_binary, args.infile) + name, bytecode = compile_solidity(solc_binary, file) - libraries[address] = bytecode + contract = ETHContract(bytecode, name=name) - print(libraries) + contracts.append(contract) -# Get encoded bytecode for main module - -if (args.code): - encoded_bytecode = args.code -elif (args.address): - encoded_bytecode = eth.eth_getCode(args.address) -elif (args.infile): - encoded_bytecode = compile_solidity(solc_binary, args.infile) - print(encoded_bytecode) else: - exitWithError("No input bytecode. Please provide EVM code via -c BYTECODE, -a ADDRESS, or -i SOLIDITY_FILE") - -logging.debug("Input bytecode: " + encoded_bytecode) + exitWithError("No input bytecode. Please provide EVM code via -c BYTECODE, -a ADDRESS, or -i SOLIDITY_FILES") -try: - disassembly = Disassembly(encoded_bytecode) -except binascii.Error: - exitWithError("Disassembler: Invalid code string.") +logging.info("Loaded " + str(len(contracts)) + " contracts") # Process commands if (args.disassemble): - easm_text = disassembly.get_easm() + easm_text = contracts[0].get_easm() sys.stdout.write(easm_text) elif (args.graph): @@ -171,7 +158,7 @@ elif (args.graph): if args.enable_physics is not None: physics = True - html = generate_callgraph(disassembly, args.enable_physics) + html = generate_callgraph(contracts[0].get_disassembly(), args.enable_physics) try: with open(args.graph, "w") as f: @@ -182,36 +169,36 @@ elif (args.graph): elif (args.xrefs): - contract = ETHContract(encoded_bytecode) + contract = ETHContract(bytecode) print("\n".join(contract.get_xrefs())) elif (args.fire_lasers): - laserfree.fire(disassembly) + laserfree.fire(contracts[0].get_disassembly()) elif (args.trace): if (args.code): - encoded_bytecode = args.code + bytecode = args.code elif (args.address): if args.ipc: eth = EthIpc() - encoded_bytecode = eth.eth_getCode(args.address) + bytecode = eth.eth_getCode(args.address) else: eth = EthJsonRpc(args.rpchost, args.rpcport, args.rpctls) - encoded_bytecode = eth.eth_getCode(args.address) + bytecode = eth.eth_getCode(args.address) else: exitWithError("Disassembler: Provide the input bytecode via -c BYTECODE or --id ID") if (args.data): - trace = evm.trace(encoded_bytecode, args.data) + trace = evm.trace(bytecode, args.data) else: - trace = evm.trace(encoded_bytecode) + trace = evm.trace(bytecode) for i in trace: if (re.match(r'^PUSH.*', i['op'])): diff --git a/mythril/ether/ethcontract.py b/mythril/ether/ethcontract.py index 456dffc3..f8aa372f 100644 --- a/mythril/ether/ethcontract.py +++ b/mythril/ether/ethcontract.py @@ -6,11 +6,11 @@ from ethereum import utils class ETHContract(persistent.Persistent): - def __init__(self, code, creation_code = ""): + def __init__(self, code, creation_code = "", name = ""): self.code = code self.creation_code = creation_code - + self.name = name def get_xrefs(self): @@ -30,7 +30,7 @@ class ETHContract(persistent.Persistent): return xrefs - def get_instruction_list(self): + def get_disassembly(self): return asm.disassemble(util.safe_decode(self.code)) diff --git a/mythril/ether/util.py b/mythril/ether/util.py index 52dc63dd..35362161 100644 --- a/mythril/ether/util.py +++ b/mythril/ether/util.py @@ -2,6 +2,8 @@ from ethereum.abi import encode_abi, encode_int from ethereum.utils import zpad from ethereum.abi import method_id import subprocess +import binascii +import os import re @@ -14,9 +16,8 @@ def safe_decode(hex_encoded_string): def compile_solidity(solc_binary, file): output = subprocess.check_output(["solc", "--bin-runtime", file], stderr=subprocess.DEVNULL) - - m = re.search(r"runtime part: \\n([0-9a-f]+)\\n", str(output)) - return m.group(1) + m = re.search(r":(.*?) =======\\nBinary of the runtime part: \\n([0-9a-f]+)\\n", str(output)) + return [m.group(1), m.group(2)] def encode_calldata(func_name, arg_types, args): @@ -24,3 +25,7 @@ def encode_calldata(func_name, arg_types, args): function_selector = zpad(encode_int(mid), 4) args = encode_abi(arg_types, args) return "0x" + function_selector.hex() + args.hex() + + +def get_random_address(): + return binascii.b2a_hex(os.urandom(20)).decode('UTF-8') diff --git a/tests/ethcontract_test.py b/tests/ethcontract_test.py index 118cebef..faf3bd25 100644 --- a/tests/ethcontract_test.py +++ b/tests/ethcontract_test.py @@ -14,7 +14,7 @@ class Getinstruction_listTestCase(ETHContractTestCase): contract = ETHContract(self.code, self.creation_code) - instruction_list = contract.get_instruction_list() + instruction_list = contract.get_disassembly() self.assertEqual(len(instruction_list), 71, 'Error disassembling code using ETHContract.get_instruction_list()')