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
from mythril.disassembler.disassembly import Disassembly
from mythril.disassembler.callgraph import generate_callgraph
from mythril.ether.contractstorage import get_persistent_storage
from mythril.ether.ethcontract import ETHContract
from mythril.rpc.client import EthJsonRpc
from mythril.ipc.client import EthIpc
from ethereum import utils
from laser.ethereum import laserfree
import logging
import binascii
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', '--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('-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')


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

args = parser.parse_args()

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


# Establish RPC/IPC connection if necessary

if (args.address or args.infile):
        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))

# Get bytecode for main module

if (args.code):
    encoded_bytecode = args.code
elif (args.address):
    encoded_bytecode = eth.eth_getCode(args.address)
elif (args.infile):
    pass
else:
    exitWithError("No input bytecode. Please provide the code via -c BYTECODE or -a address")


logging.debug("Input bytecode: " + encoded_bytecode)


try:
    disassembly = Disassembly(encoded_bytecode)
except binascii.Error:
    exitWithError("Disassembler: Invalid code string.")

# Process commands

if (args.disassemble):

    easm_text = disassembly.get_easm()
    sys.stdout.write(easm_text)

elif (args.graph):

    if args.enable_physics is not None:
        physics = True 

    html = generate_callgraph(disassembly, 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(encoded_bytecode)

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

elif (args.fire_lasers):

    laserfree.fire(disassembly)

elif (args.trace):

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

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

        else:
            eth = EthJsonRpc(args.rpchost, args.rpcport, args.rpctls)
            encoded_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)

    else:
        trace = evm.trace(encoded_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'])

elif 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):
        if args.ipc:
                contract_storage.initialize(args.rpchost, args.rpcport, args.rpctls, args.sync_all, args.ipc)
        else:
            contract_storage.initialize(args.rpchost, args.rpcport, args.rpctls, args.sync_all, args.ipc)

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

else:
    parser.print_help()