from pwn import *

# Context setup
context.log_level = 'error' # Reduce noise, we only care about the flag or retrying

# The Payload script to send to the binary
PAYLOAD = """
import sys
target = -1
next_box = -1
while True:
    try:
        line = sys.stdin.readline()
        if not line: break
        line = line.strip()
        if "You are player number" in line:
            parts = line.split()
            current_player = int(parts[4].replace('.', ''))
            if target == -1:
                target = current_player
                next_box = current_player
            print(next_box)
            sys.stdout.flush()
        elif line.startswith("FOUND"):
            parts = line.split()
            next_box = int(parts[1])
    except: break
"""

def solve():
    # Replace './challenge' with the actual binary name
    # If remote, use: p = remote('IP', PORT)
    # p = remote('chall.polygl0ts.ch', 6667)
    p = process(['./chal'], env={'FLAG': 'CTF{test_flag}'}) 

    print("[*] Attempting to solve 100 Prisoners...")
    
    for i in range(1, 101):
        try:
            # Wait for prompt "Provide Python script for player X"
            p.recvuntil(b'EOF', timeout=1) 
            
            # Send our payload + EOF
            p.sendline(PAYLOAD.encode())
            p.sendline(b"EOF")
            
            # Check if this player failed
            # The binary prints "Player X failed" or "Player X succeeded"
            response = p.recvline_pred(lambda x: b"succeeded" in x or b"failed" in x)
            
            if b"failed" in response:
                p.close()
                return False
                
        except EOFError:
            p.close()
            return False

    # If we exit the loop, everyone succeeded!
    print("[+] Success! Probability gods smiled upon us.")
    p.interactive() # Get the flag
    return True

attempt = 1
while True:
    print(f"[-] Run #{attempt}")
    if solve():
        break
    attempt += 1