main.py

# main.py
import os
import uuid
import shutil
import subprocess
import tempfile
import time
from typing import List, Optional

from fastapi import FastAPI, HTTPException, Request
from pydantic import BaseModel

API_KEY = os.environ.get("API_KEY")  # set this in Render env
MAX_RUNTIME_SECONDS = int(os.environ.get("MAX_RUNTIME_SECONDS", "10"))
MAX_OUTPUT_BYTES = int(os.environ.get("MAX_OUTPUT_BYTES", "200_000"))  # limit output size

app = FastAPI(title="Py Runner API (CPython)", version="1.0")


class RunRequest(BaseModel):
    code: str
    install: Optional[List[str]] = []  # list of pip packages to install (optional)


def safe_run_python_script(script_path: str, timeout: int):
    """
    Run `python script_path` in subprocess with timeout and capture stdout/stderr.
    Returns tuple (returncode, stdout, stderr).
    """
    try:
        proc = subprocess.run(
            ["python3", script_path],
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        stdout = proc.stdout[:MAX_OUTPUT_BYTES]
        stderr = proc.stderr[:MAX_OUTPUT_BYTES]
        return proc.returncode, stdout, stderr
    except subprocess.TimeoutExpired as e:
        # kill proc handled by subprocess
        return -1, e.stdout or "", f"Timeout expired after {timeout} seconds"
    except Exception as e:
        return -2, "", f"Execution error: {e}"


@app.post("/run")
async def run_code(req: Request, payload: RunRequest):
    # API key validation
    header = req.headers.get("Authorization") or req.headers.get("X-Api-Key")
    if API_KEY:
        if not header:
            raise HTTPException(status_code=401, detail="Missing API key")
        # allow "Bearer <key>" or plain key
        if header.startswith("Bearer "):
            key_val = header.split(" ", 1)[1].strip()
        else:
            key_val = header.strip()
        if key_val != API_KEY:
            raise HTTPException(status_code=403, detail="Invalid API key")

    code = payload.code
    install = payload.install or []

    # create ephemeral working dir
    workdir = tempfile.mkdtemp(prefix="py-run-")
    try:
        # optional: create venv per request (expensive). We'll use global python env:
        # Install packages if requested (careful: heavy and slow)
        pip_install_logs = []
        if install:
            # safeguard: allow only a whitelist pattern? (OPTIONAL)
            for pkg in install:
                # basic sanitization
                if any(c in pkg for c in [";", "&", "|", "$"]):
                    raise HTTPException(status_code=400, detail=f"Invalid package name: {pkg}")
                # run pip install
                p = subprocess.run(
                    ["pip", "install", pkg],
                    capture_output=True,
                    text=True,
                    cwd=workdir,
                    timeout=300,
                )
                pip_install_logs.append({"pkg": pkg, "rc": p.returncode, "stdout": p.stdout[:1000], "stderr": p.stderr[:1000]})
                if p.returncode != 0:
                    return {"error": f"pip install failed for {pkg}", "pip": pip_install_logs}

        # write code to temp file
        script_path = os.path.join(workdir, "script.py")
        with open(script_path, "w", encoding="utf-8") as f:
            f.write(code)

        # run script safely
        start = time.time()
        rc, stdout, stderr = safe_run_python_script(script_path, timeout=MAX_RUNTIME_SECONDS)
        duration = time.time() - start

        return {
            "return_code": rc,
            "stdout": stdout,
            "stderr": stderr,
            "duration_seconds": round(duration, 3),
            "pip_logs": pip_install_logs,
        }
    finally:
        # cleanup
        try:
            shutil.rmtree(workdir)
        except Exception:
            pass