#!/Users/berndt/.pyenv/versions/3.5.2/bin/python
"""mythril.py: Bug hunting on the Ethereum blockchain

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

from ether import asm,evm,util
from ether.contractstorage import get_persistent_storage
from ethereum import utils
import sys
import argparse
from rpc.client import EthJsonRpc

import os


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, use 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', default='0x0123456789ABCDEF0123456789ABCDEF01234567', help='contract address')
parser.add_argument('-o', '--outfile')
parser.add_argument('--data', help='message call input data for tracing')
parser.add_argument('--search', help='search the contract database')
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')


app_root = os.path.dirname(os.path.realpath(__file__))
db_dir = os.path.join(app_root, "database")

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))
    elif (args.infile):

        try:

            encoded_bytecode = util.file_to_string(args.infile).rstrip()

        except Exception as e:
            exitWithError("Exception loading bytecode from file: " + str(e))

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

    disassembly = asm.disassemble(util.safe_decode(encoded_bytecode))

    easm_text = asm.disassembly_to_easm(disassembly)

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

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):
        output = evm.trace(util.safe_decode(encoded_bytecode), args.address, args.data)

    else:
        output = evm.trace(util.safe_decode(encoded_bytecode), args.address)

    print(output)

elif (args.search):

    contract_storage.search(args.search, searchCallback)

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

elif (args.hash):
    print(utils.sha3(args.hash)[:4].hex())

else:
    parser.print_help()
