mythril/myth

#!/usr/bin/env python
"""mythril.py: Bug hunting on the Ethereum blockchain

   http://www.github.com/b-mueller/mythril
"""

from mythril.ether import evm, util
from mythril.disassembler.callgraph import generate_callgraph
from mythril.ether.contractstorage import get_persistent_storage
from mythril.ether.ethcontract import ETHContract
from mythril.ether.util import compile_solidity
from mythril.rpc.client import EthJsonRpc
from mythril.ipc.client import EthIpc
from ethereum import utils
from laser.ethereum import laserfree
import logging
import sys
import argparse
import os
import re


def searchCallback(code_hash, code, addresses, balances):
    print("Matched contract with code hash " + code_hash)

    for i in range(0, len(addresses)):
        print("Address: " + addresses[i] + ", balance: " + str(balances[i]))


def exitWithError(message):
    print(message)
    sys.exit()


parser = argparse.ArgumentParser(description='Bug hunting on the Ethereum blockchain')

commands = parser.add_argument_group('commands')
commands.add_argument('-d', '--disassemble',  action='store_true', help='disassemble')
commands.add_argument('-g', '--graph', help='generate a control flow graph', metavar='OUTPUT_FILE')
commands.add_argument('-x', '--fire-lasers', action='store_true', help='detect vulnerabilities')
commands.add_argument('-t', '--trace', action='store_true', help='trace, use with -c or -a and --data (optional)')
commands.add_argument('-s', '--search', help='search the contract database')
commands.add_argument('--xrefs', action='store_true', help='get xrefs from a contract, use with -c or -a')
commands.add_argument('--hash', help='calculate function signature hash', metavar='SIGNATURE')
commands.add_argument('--init-db', action='store_true', help='initialize the contract database')

inputs = parser.add_argument_group('input arguments')
inputs.add_argument('-c', '--code', help='hex-encoded bytecode string ("6060604052...")', metavar='BYTECODE')
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')

options = parser.add_argument_group('options')
options.add_argument('--sync-all', action='store_true', help='Also sync contracts with zero balance')
options.add_argument('--rpchost', default='127.0.0.1', help='RPC host')
options.add_argument('--rpcport', type=int, default=8545, help='RPC port')
options.add_argument('--rpctls', type=bool, default=False, help='RPC port')
options.add_argument('--ipc', help='use IPC interface instead of RPC', action='store_true')
options.add_argument('--enable-physics', type=bool, default=False, help='enable graph physics simulation')
options.add_argument('-v', type=int, help='log level (0-2)', metavar='LOG_LEVEL')

# Get config values

try:
    db_dir = os.environ['DB_DIR']
except KeyError:
    db_dir = None

try:
    solc_binary = os.environ['SOLC']
except KeyError:
    solc_binary = 'solc'

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):
    parser.print_help()
    sys.exit()

if (args.v):
    if (0 <= args.v < 3):
        logging.basicConfig(level=[logging.NOTSET, logging.INFO, logging.DEBUG][args.v])

elif (args.hash):
    print("0x" + utils.sha3(args.hash)[:4].hex())
    sys.exit()

# Database search ops

if args.search or args.init_db:

    contract_storage = get_persistent_storage(db_dir)

    if (args.search):

        try:
            contract_storage.search(args.search, searchCallback)
        except SyntaxError:
            exitWithError("Syntax error in search expression.")

    elif (args.init_db):
            contract_storage.initialize(args.rpchost, args.rpcport, args.rpctls, args.sync_all, args.ipc)

    sys.exit()

# Establish RPC/IPC connection if necessary

if (args.address or args.infiles):
        if args.ipc:
            try:
                eth = EthIpc()

            except Exception as e:
                exitWithError("Error establishing IPC connection: " + str(e))
        else:
            try:
                eth = EthJsonRpc(args.rpchost, args.rpcport, args.rpctls)

            except Exception as e:
                exitWithError("Error establishing RPC connection: " + str(e))

# Load input contracts

contracts = []

if (args.code):
    contracts.append(ETHContract(args.code, address = util.get_random_address()))
elif (args.address):
    contracts.append(ETHContract(eth.eth_getCode(args.address), address = args.address))
elif (args.infiles):
    files = args.infiles.split(",")

    counter = 0

    for file in files:

        name, bytecode = compile_solidity(solc_binary, file)

        # Max. 16 contracts supported

        contract = ETHContract(bytecode, name = name, address = "0x" + hex(counter)[2:] * 40)
        counter += 1

        contracts.append(contract)
        logging.info(contract.name + " at " + contract.address)

else:
    exitWithError("No input bytecode. Please provide EVM code via -c BYTECODE, -a ADDRESS, or -i SOLIDITY_FILES")

# Commands

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:
        print(str(contract.as_dict()))
        modules.append(contract.as_dict())

    html = generate_callgraph(modules, args.enable_physics)

    try:
        with open(args.graph, "w") as f:
            f.write(html)
    except Exception as e:

        print("Error saving graph: " + str(e))

elif (args.xrefs):

    contract = ETHContract(bytecode)

    print("\n".join(contract.get_xrefs()))

elif (args.fire_lasers):

    modules = []

    for contract in contracts:
        modules.append(contract.as_dict())

    laserfree.fire(modules)

elif (args.trace):

    if (args.code):
        bytecode = args.code

    elif (args.address):
        if args.ipc:
            eth = EthIpc()
            bytecode = eth.eth_getCode(args.address)

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

    if (args.data):
        trace = evm.trace(bytecode, args.data)

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


else:
    parser.print_help()