#!/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.disassembly import Disassembly
from mythril.disassembler.callgraph import generate_callgraph
from mythril.ether.contractstorage import get_persistent_storage
from mythril.rpc.client import EthJsonRpc
from ethereum import utils
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')

parser.add_argument('-d', '--disassemble',  action='store_true', help='disassemble, specify input with -c or -a')
parser.add_argument('-t', '--trace', action='store_true', help='trace, use with -c or -a and --data (optional)')
parser.add_argument('-c', '--code', help='hex-encoded bytecode string ("6060604052...")', metavar='BYTECODE')
parser.add_argument('-a', '--address', help='contract address')
parser.add_argument('-o', '--outfile')
parser.add_argument('-g', '--graph', help='when disassembling, also generate a callgraph', metavar='OUTPUT_FILE')
parser.add_argument('--data', help='message call input data for tracing')
parser.add_argument('--search', help='search the contract database')
parser.add_argument('--xrefs', help='get xrefs from contract in database', metavar='CONTRACT_HASH')
parser.add_argument('--hash', help='calculate function signature hash', metavar='SIGNATURE')
parser.add_argument('--init-db', action='store_true', help='Initialize the contract database')
parser.add_argument('--sync-all', action='store_true', help='Also sync contracts with zero balance')
parser.add_argument('--rpchost', default='127.0.0.1', help='RPC host')
parser.add_argument('--rpcport', type=int, default=8545, help='RPC port')

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

contract_storage = get_persistent_storage(db_dir)

args = parser.parse_args()

if (args.disassemble):

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

        try:
            eth = EthJsonRpc(args.rpchost, args.rpcport)

            encoded_bytecode = eth.eth_getCode(args.address)

        except Exception as e:
            exitWithError("Exception loading bytecode via RPC: " + str(e))

    else:
        exitWithError("Disassembler: Provide the input bytecode via -c BYTECODE or --id ID")

    try:
        disassembly = Disassembly(encoded_bytecode)
        # instruction_list = asm.disassemble(util.safe_decode(encoded_bytecode))
    except binascii.Error:
        exitWithError("Disassembler: Invalid code string.")

    easm_text = disassembly.get_easm()

    if (args.outfile):
        util.string_to_file(args.outfile, easm_text)
    else:
        sys.stdout.write(easm_text)

    if (args.graph):
        generate_callgraph(disassembly, args.graph)


elif (args.trace):

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

    elif (args.address):

        eth = EthJsonRpc(args.rpchost, args.rpcport)

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

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

elif (args.xrefs):

    try:
        contract_hash = util.safe_decode(args.xrefs)
    except binascii.Error:
        exitWithError("Invalid contract hash.")

    try:
        contract = contract_storage.get_contract_by_hash(contract_hash)
        print("\n".join(contract.get_xrefs()))
    except KeyError:
        exitWithError("Contract not found in the database.")        

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

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

else:
    parser.print_help()