#!/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.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 mythril.support.loader import DynLoader
from mythril.exceptions import CompilerError
from mythril.analysis.symbolic import StateSpace
from mythril.analysis.callgraph import generate_graph
from mythril.analysis.security import fire_lasers
from ethereum import utils
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('-l', '--dynld', action='store_true', help='auto-load dependencies (experimental)')
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'

# Parse cmdline args

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 / compile input contracts

contracts = []

if (args.code):
    contracts.append(ETHContract(args.code, name="MAIN", address = util.get_indexed_address(0)))
elif (args.address):
    contracts.append(ETHContract(eth.eth_getCode(args.address), name=args.address, address = args.address))
elif (len(args.solidity_file)):
    index = 0

    for file in args.solidity_file:

        file = file.replace("~", str(Path.home())) # Expand user path

        try:
            name, bytecode = compile_solidity(solc_binary, file)
        except CompilerError as e:
            exitWithError(e)

        # Max. 16 input files supported!

        contract = ETHContract(bytecode, name = name, address = util.get_indexed_address(index))
        index += 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.trace):

    if (args.data):
        trace = evm.trace(contracts[0].code, args.data)

    else:
        trace = evm.trace(contracts[0].code)

    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.xrefs):

    print("\n".join(contracts[0].get_xrefs()))

elif (args.graph) or (args.fire_lasers):

    if (args.graph):

        if (args.dynld):
            states = StateSpace(contracts, dynamic_loader=DynLoader(eth), simplified=True)
        else:
            states = StateSpace(contracts, simplified=True)

        if args.enable_physics is not None:
            physics = True

        html = generate_graph(states, args.enable_physics)

        try:
            with open(args.graph, "w") as f:
                f.write(html)
        except Exception as e:
            print("Error saving graph: " + str(e))

    else:

        if (args.dynld):
            states = StateSpace(contracts, dynamic_loader=DynLoader(eth))
        else:
            states = StateSpace(contracts)

        fire_lasers(states)

else:
    parser.print_help()