Load contracts into list

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'])):

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))
 

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')

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()')