117 lines
2.8 KiB
Python
117 lines
2.8 KiB
Python
#!/usr/bin/python2
|
|
|
|
import argparse
|
|
import re
|
|
import sys
|
|
|
|
from . import Junk
|
|
from .AsciiPainter import AsciiPainter
|
|
from .Emulator import REGISTERS, Emulator
|
|
from .TikzPainter import TikzPainter
|
|
|
|
ARG_DESC = """
|
|
A silly x86-64 emulator and stack visualizer.
|
|
"""
|
|
|
|
PAINTER_ID_TO_CLASS = {'ascii': AsciiPainter, 'tikz': TikzPainter}
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description=ARG_DESC)
|
|
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('call 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()
|