NoN-CTF/Misc/Medium/Grinch's Lottery/main.py

#!/usr/bin/env python3
import os
import sys
import socketserver
import random
import io

FLAG_PRICE = 1_000_000
INITIAL_COINS = 50  # tiny starting balance
MAX_TURNS = 10
MAX_LEGAL_LOAN = 1000  # max positive amount per loan


def load_flag():
    # Flag is supplied via environment variable
    return os.getenv("FLAG", "XMAS{dummy_flag_for_testing}")


def to_int(s: str) -> int:
    return int(s.strip())


def add_coins(balance: int, delta: int) -> int:
    """
    Update the balance using 32-bit unsigned wraparound.

    This is the core bug: using a 32-bit "safe" ledger but letting the
    user control delta directly (including large negatives).
    """
    return (balance + delta) & 0xFFFFFFFF


def play_game(reader, writer):
    """
    Core game logic. Works with any text reader/writer pair.
    - CLI mode: stdin/stdout
    - TCP mode: wrapped socket rfile/wfile
    """
    flag = load_flag()
    coins = INITIAL_COINS
    price = FLAG_PRICE
    turns = 0

    def send(line: str = ""):
        writer.write(line + "\n")
        writer.flush()

    send("=== Grinch's Lottery ===")
    send("Welcome to Whoville's shadiest Christmas lottery!")
    send("You start with a few coal-coins. Win enough to buy the special prize.")
    send(f"The shiny present costs {price} coins.")
    send("")
    while turns < MAX_TURNS:
        turns += 1
        send(f"\n[Turn {turns}/{MAX_TURNS}] Your balance: {coins} coins")
        send("Choose an action:")
        send("  1) Work a little (earn some coins)")
        send("  2) Buy lottery tickets (10 coins each)")
        send("  3) Take a shady loan from the Grinch")
        send("  4) Buy the shiny present")
        send("  5) Exit")
        send("> ")
        choice = reader.readline()
        if not choice:
            break
        choice = choice.strip()

        if choice == "1":
            # Small honest income, no way to reach the flag fairly
            earn = random.randint(1, 20)
            coins = add_coins(coins, earn)
            send(f"You did some odd jobs for the Whos and earned {earn} coins.")

        elif choice == "2":
            # Lottery tickets that cost coins and rarely pay out
            send("How many tickets do you want to buy? (10 coins each)")
            send("> ")
            line = reader.readline()
            if not line:
                break
            try:
                n = to_int(line)
            except ValueError:
                send("That's not even a number. The Grinch laughs you out of town.")
                break

            # Make lottery 'honest' so it's not the exploit:
            if n <= 0:
                send("You must buy at least 1 ticket, no refunds.")
                continue

            cost = 10 * n

            # If you can't afford it, don't touch the balance at all
            if cost > coins:
                send("You don't have enough coins for that many tickets.")
                send("The Grinch refuses to give you credit for lottery tickets.")
                continue

            # Now this subtraction is safe (no wrap) because cost <= coins
            coins = add_coins(coins, -cost)
            send(f"You bought {n} tickets for {cost} coins.")

            wins = 0
            # Cap iteration count so people can't hang the server
            for _ in range(min(n, 1000)):
                if random.random() < 0.05:
                    prize = random.randint(5, 50)
                    coins = add_coins(coins, prize)
                    wins += 1
            if wins:
                send(f"You won on {wins} ticket(s)! Your luck slightly improves.")
            else:
                send("No luck this time. The Grinch grins.")

        elif choice == "3":
            # Shady loan service with a "safety limit"... that only checks positive amounts.
            send("How much do you want from the Grinch's loan shark service?")
            send("> ")
            line = reader.readline()
            if not line:
                break
            try:
                delta = to_int(line)
            except ValueError:
                send("You mumble nonsense. The Grinch gets bored and leaves.")
                break

            # Safety check that *only* constrains positive loans.
            if delta > MAX_LEGAL_LOAN:
                send("Too risky! The Grinch's auditor caps loans at "
                     f"{MAX_LEGAL_LOAN} coins per request.")
                send("Smaller loans only, please.")
                continue

            # Negative values are not checked at all. This is the intended bug:
            # a large negative 'repayment' will underflow and wrap to a huge balance.
            before = coins
            coins = add_coins(coins, delta)
            send(f"Ledger update: {before} -> {coins} (delta: {delta})")

        elif choice == "4":
            # Check if the player can afford the flag
            if coins >= price:
                send("You jingle enough coins to buy the shiny present!")
                send(f"Inside you find a note: {flag}")
                send("Merry XMAS and beware of greedy Grinches.")
                return
            else:
                price += FLAG_PRICE  # price goes up if you try to buy without enough coins
                send(f"Not enough coins. The Grinch cackles and raises the price to {price} coins.")

        elif choice == "5":
            send("You leave the lottery, slightly colder but maybe wiser.")
            return

        else:
            send("The Grinch doesn't recognize that option and kicks you out.")
            return

    send("The Grinch closes the stand for the night. Come back with a better plan.")


# === TCP server glue (thread-per-connection, like gunicorn workers) ===

class GrinchHandler(socketserver.StreamRequestHandler):
    def handle(self):
        # Simple timeout so connections don't hang forever
        self.request.settimeout(60)

        # Wrap the raw socket file descriptors into text-mode file objects
        reader = io.TextIOWrapper(self.rfile, encoding="utf-8", newline="\n")
        writer = io.TextIOWrapper(
            self.wfile,
            encoding="utf-8",
            newline="\n",
            write_through=True,
        )

        try:
            play_game(reader, writer)
        except Exception:
            # Hide internals from players
            try:
                writer.write("An error occurred. The Grinch eats the stack trace.\n")
                writer.flush()
            except Exception:
                pass


class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    # Basic "load balancing": one thread per client connection
    allow_reuse_address = True
    daemon_threads = True


def run_server(host: str = "0.0.0.0", port: int = 9000):
    with ThreadedTCPServer((host, port), GrinchHandler) as server:
        server.serve_forever()


if __name__ == "__main__":
    # Usage:
    #   ./grinch_lottery.py             -> CLI mode on stdin/stdout
    #   ./grinch_lottery.py server      -> TCP server 0.0.0.0:9000
    #   ./grinch_lottery.py server HOST PORT
    if len(sys.argv) >= 2 and sys.argv[1] == "server":
        host = sys.argv[2] if len(sys.argv) >= 3 else "0.0.0.0"
        port = int(sys.argv[3]) if len(sys.argv) >= 4 else 9000
        run_server(host, port)
    else:
        # CLI mode: interact via stdin/stdout (no network)
        reader = sys.stdin
        writer = sys.stdout
        play_game(reader, writer)