#!/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
from pathlib import Path
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')
parser.add_argument("solidity_file", nargs='*')


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 contract, use with --data (optional)')
commands.add_argument('-s', '--search', help='search the contract database', metavar='EXPRESSION')
commands.add_argument('--xrefs', action='store_true', help='get xrefs from a contract')
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('-a', '--address', help='pull contract from the blockchain', 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('--infura-mainnet', action='store_true', help='Use Infura Node service, equivalent to: --rpchost=mainnet.infura.io --rpcport=443 --rpctls="True"')
options.add_argument('--infura-rinkeby', action='store_true', help='Use Infura Node service, equivalent to: --rpchost=rinkeby.infura.io --rpcport=443 --rpctls="True"')
options.add_argument('--infura-kovan', action='store_true', help='Use Infura Node service, equivalent to: --rpchost=kovan.infura.io --rpcport=443 --rpctls="True"')
options.add_argument('--infura-ropsten', action='store_true', help='Use Infura Node service, equivalent to: --rpchost=ropsten.infura.io --rpcport=443 --rpctls="True"')
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 len(args.solidity_file)):
        if args.ipc:
            try:
                eth = EthIpc()

            except Exception as e:
                exitWithError("Error establishing IPC connection: " + str(e))
        else:
            try:
                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)

            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 (len(args.solidity_file)):
    counter = 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

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

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