#!/usr/bin/env python3
"""
Sandbox Slip - CTF challenge (medium, misc)

Supports two modes:
 - TCP (default): listens on HOST:PORT and serves a netcat-style REPL for each connection.
 - STDIN: reads lines from stdin and writes results to stdout (local interactive mode).

Usage:
  python3 server.py            # TCP mode (default)
  python3 server.py --mode stdin   # stdin/stdout REPL
  python3 server.py --host 0.0.0.0 --port 9000
"""

import socketserver
import threading
import sys
import traceback
import os
import argparse

HOST = "0.0.0.0"
PORT = 9000
FLAG_FILE = "flag.txt"
WELCOME = """Welcome to PySafe REPL v0.3
This REPL runs your Python expressions in a restricted environment.
Type a single expression per line. Type QUIT to exit.
Available helper object: helper
>> """

# Vulnerable helper intentionally included to demonstrate sandbox escape.
def helper():
    """A small helper function that does nothing useful for players."""
    return "helpful"

# Ensure flag exists
if not os.path.exists(FLAG_FILE):
    with open(FLAG_FILE, "w") as f:
        f.write("CTF{sandbox_slip_helper_globals}\n")

# Prepare the sandbox environment (shared by TCP and STDIN modes)
# Note: __builtins__ is set to empty dict to simulate a naive sandbox.
safe_globals = {"__builtins__": {}}
safe_locals = {"helper": helper}

def evaluate_source(src: str, g: dict, l: dict):
    """
    Evaluate a single-line source in the restricted environment.
    Tries eval() first, then exec() for statements.
    Returns a tuple (success:bool, output:str).
    """
    src = src.rstrip("\n")
    if src.strip().upper() == "QUIT":
        return True, "QUIT"
    if src.strip() == "":
        return True, ""  # empty line -> no output
    try:
        # Try eval for expressions
        result = eval(src, g, l)
        return True, repr(result)
    except SyntaxError:
        # Try exec for statements
        try:
            exec(src, g, l)
            return True, "OK"
        except Exception as e:
            tb = traceback.format_exc()
            return False, f"Execution error: {e}"
    except Exception as e:
        return False, f"Error: {e}"

# ---------- TCP handler ----------
class REPLHandler(socketserver.StreamRequestHandler):
    def handle(self):
        addr = self.client_address[0]
        print(f"[+] Connection from {addr}")
        # Use instance-local copies so per-connection state can diverge if needed
        g = safe_globals
        l = safe_locals
        try:
            self.wfile.write(WELCOME.encode())
            while True:
                line = self.rfile.readline()
                if not line:
                    break
                src = line.decode().rstrip("\n")
                if src.strip().upper() == "QUIT":
                    self.wfile.write(b"Goodbye.\n")
                    break
                ok, out = evaluate_source(src, g, l)
                # If source was empty, don't print extra blank line (just prompt)
                if out == "":
                    self.wfile.write(b">> ")
                    continue
                # If QUIT sentinel returned
                if out == "QUIT":
                    self.wfile.write(b"Goodbye.\n")
                    break
                # Send result or error
                self.wfile.write((out + "\n").encode())
                self.wfile.write(b">> ")
        except ConnectionResetError:
            pass
        except Exception as e:
            print("Handler exception:", e)
        finally:
            print(f"[-] Disconnected {addr}")

class ThreadedTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True

def run_tcp_server(host: str, port: int):
    srv = ThreadedTCPServer((host, port), REPLHandler)
    print(f"Sandbox Slip (TCP) listening on {host}:{port}")
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        print("Shutting down TCP server...")
        srv.shutdown()
        srv.server_close()

# ---------- STDIN / STDOUT loop ----------
def run_stdin_server():
    """
    Interactive REPL on stdin/stdout. Useful for local testing.
    Reads from sys.stdin, writes responses to sys.stdout, prompts with '>> '.
    """
    g = safe_globals
    l = safe_locals
    sys.stdout.write(WELCOME)
    sys.stdout.flush()
    try:
        while True:
            # Read one line from stdin
            line = sys.stdin.readline()
            if not line:
                # EOF
                break
            src = line.rstrip("\n")
            if src.strip().upper() == "QUIT":
                sys.stdout.write("Goodbye.\n")
                sys.stdout.flush()
                break
            ok, out = evaluate_source(src, g, l)
            if out == "":
                sys.stdout.write(">> ")
                sys.stdout.flush()
                continue
            if out == "QUIT":
                sys.stdout.write("Goodbye.\n")
                sys.stdout.flush()
                break
            sys.stdout.write(out + "\n")
            sys.stdout.write(">> ")
            sys.stdout.flush()
    except KeyboardInterrupt:
        sys.stdout.write("\nInterrupted. Exiting.\n")
        sys.stdout.flush()

# ---------- CLI ----------
def main():
    parser = argparse.ArgumentParser(description="Sandbox Slip (TCP or STDIN REPL)")
    parser.add_argument("--mode", choices=["tcp", "stdin"], default="tcp",
                        help="Run mode: 'tcp' to listen on network (default), 'stdin' for local REPL")
    parser.add_argument("--host", default=HOST, help="Host to bind in TCP mode")
    parser.add_argument("--port", type=int, default=PORT, help="Port to bind in TCP mode")
    args = parser.parse_args()

    if args.mode == "stdin":
        run_stdin_server()
    else:
        run_tcp_server(args.host, args.port)

if __name__ == "__main__":
    main()