#!/usr/bin/env python
"""mythril.py: Ethereum VM bytecode assembler/ disassembler

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

from ether import asm,evm,util
import sys
import codecs
import argparse


def exitWithError(message):
    print(message)
    sys.exit()


parser = argparse.ArgumentParser(description='Ethereum VM bytecode assembler/ disassembler')

parser.add_argument('-d', '--disassemble',  action='store_true', help='disassemble, use with -c, -f or --txid')
parser.add_argument('-a', '--assemble', help='produce bytecode from easm input file', metavar='INPUTFILE')
parser.add_argument('-t', '--trace',  action='store_true', help='trace bytecode provided via the -c or -f argument')
parser.add_argument('-c', '--code', help='hex-encoded bytecode string ("6060604052...")', metavar='BYTECODE')
parser.add_argument('-o', '--outfile')
parser.add_argument('-f', '--infile', metavar='INPUTFILE')
parser.add_argument('--txid', help='id of contract creation transaction')
parser.add_argument('--rpchost', default='127.0.0.1', help='RPC host')
parser.add_argument('--rpcport', type=int, default=8545, help='RPC port')

args = parser.parse_args()


if (args.disassemble):

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

        try:

            encoded_bytecode = util.bytecode_from_blockchain(args.txid, args.rpchost, args.rpcport)

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

    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.message))

    else:
        exitWithError("Disassembler: Provide the input bytecode via -c BYTECODE, -f INPUT_FILE or --txid TXID")

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

    easm = util.file_to_string(args.assemble)

    disassembly = asm.easm_to_disassembly(easm)

    assembly = asm.assemble(disassembly)

    if (args.outfile):
        util.raw_bytes_to_file(args.outfile, assembly)
    else:
        print("0x" + str(codecs.encode(assembly, "hex_codec"), 'ascii'))

elif (args.trace):

    if args.code:     

        bytecode = util.safe_decode(args.code)

    elif (args.infile):

        try:

            bytecode = util.file_to_raw_bytes(args.infile)

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

    else:
        exitWithError("Trace: Provide the input bytecode using -c BYTECODE or -f INPUT_FILE")

    evm.trace(bytecode)

else:
    parser.print_help()
