infernal-interpreter/main.py

103 lines
2.9 KiB
Python
Executable File

#!/usr/bin/python2
ARG_DESC = '''
A silly x86-64 emulator and stack visualizer.
'''
LICENSE = '''
This program is BEER-WARE.
<jonjmaa@gmail.com> 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("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()