infernal-interpreter/infernal_interpreter/__main__.py

#!/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()