#!/usr/bin/python2

ARG_DESC = '''
A silly x86-64 emulator and stack visualizer.
'''

LICENSE = '''
This program is BEER-WARE.

Jmaa wrote this program.  As long as you retain this notice you
can do whatever you want with this stuff. If we meet some day, and you think
this stuff is worth it, you can buy me a beer in return.
'''

import sys

import re
import argparse

from Emulator import Emulator, CodeParseException, REGISTERS
from TikzPainter import TikzPainter
from AsciiPainter import AsciiPainter
import Junk

PAINTER_ID_TO_CLASS = {
    'ascii': AsciiPainter,
    'tikz': TikzPainter
}

def parse_args ():
	parser = argparse.ArgumentParser(description = ARG_DESC, epilog = LICENSE)
	parser.add_argument('-i', '--input-file',  default='-',     help = 'File to use as input',  dest = 'input-file')
	parser.add_argument('-o', '--output-file', default='-',     help = 'File to use as output', dest = 'output-file')
	parser.add_argument('-p', '--painter',     default='ascii', help = 'Drawer to use',         dest = 'painter-name', choices=PAINTER_ID_TO_CLASS.keys())
	for register in REGISTERS:
		parser.add_argument('--'+register[1:], nargs = 1, type = int, default = None, dest = register, help = 'the initial value of the register')
	args = vars(parser.parse_args())

	# Determine args #
	registers_init = {}
	for register in REGISTERS:
            registers_init[register[1:]] = Junk.Junk('old '+register[1:])
	#
	registers_init['rip'] = 0
	registers_init['rsp'] = 0
	registers_init['rbp'] = Junk.Junk('old bp')
	#
	for register in REGISTERS:
		if args[register]:
			registers_init[register[1:]] = args[register][0]

	#
	program = ""
	if args['input-file'] == '-':
		program = sys.stdin.read()
	else:
		with open(args['input-file']) as file:
			program = file.read()

	# Output file:
	output_file = None
	if args['output-file'] == '-':
		output_file = sys.stdout
	else:
		output_file = open(args['output-file'], 'w')


        # Determine painter class
        painter_class = PAINTER_ID_TO_CLASS[args['painter-name']]

	## Return
	return (program, registers_init, output_file, painter_class)

def main ():

	(program, registers_init, output_file, painter_class) = parse_args()

	# Determine registers to display
	registers_to_draw = ["%rip","%rbp","%rsp", ""]
	for match in re.findall(r"%r[a-z0-9]{1,2}", program):
		if match not in registers_to_draw:
			registers_to_draw.append(match)

	# Setup emulator and drawer
	emu = Emulator(program)
	emu.setRegs(**registers_init)
	emu.setStack(Junk.Junk("junk..."), Junk.Junk("calling eip"))

	painter = painter_class(registers_to_draw)

	# Iteratively draw states.
	painter.saveState(emu)
	for line_nr in emu:
		painter.saveState(emu, line_nr)

	# Display result
	output_file.write(painter.to_string(emu))
	output_file.close()

if __name__ == "__main__":
	main()