atlas-iac/dockerfiles/hermes-chat-sandbox-server.py

"""Small credential-free Python execution service for Hermes chat tenants."""

from __future__ import annotations

import json
import os
import resource
import signal
import subprocess
import tempfile
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

HOST = "0.0.0.0"
PORT = 9080
MAX_REQUEST_BYTES = 160 * 1024
MAX_CODE_BYTES = 128 * 1024
MAX_OUTPUT_BYTES = 100 * 1024
WORKSPACE = Path("/workspace")


def _child_limits() -> None:
    """Apply conservative CPU, memory, process, file, and descriptor limits."""
    os.setsid()
    resource.setrlimit(resource.RLIMIT_CPU, (35, 35))
    resource.setrlimit(resource.RLIMIT_AS, (768 * 1024 * 1024,) * 2)
    resource.setrlimit(resource.RLIMIT_NPROC, (32, 32))
    resource.setrlimit(resource.RLIMIT_FSIZE, (32 * 1024 * 1024,) * 2)
    resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))


def _execute(code: str) -> dict[str, object]:
    """Run one isolated Python subprocess and return bounded output."""
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", prefix="hermes-", dir="/tmp", delete=False
    ) as script:
        script.write(code)
        script_path = script.name

    env = {
        "HOME": str(WORKSPACE),
        "LANG": "C.UTF-8",
        "LC_ALL": "C.UTF-8",
        "PATH": "/usr/local/bin:/usr/bin:/bin",
        "PYTHONDONTWRITEBYTECODE": "1",
    }
    try:
        process = subprocess.Popen(
            ["python", "-I", "-B", script_path],
            cwd=WORKSPACE,
            env=env,
            stdin=subprocess.DEVNULL,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            preexec_fn=_child_limits,
        )
        try:
            stdout, stderr = process.communicate(timeout=45)
            timed_out = False
        except subprocess.TimeoutExpired:
            timed_out = True
            os.killpg(process.pid, signal.SIGKILL)
            stdout, stderr = process.communicate()
        return {
            "success": process.returncode == 0 and not timed_out,
            "exit_code": process.returncode,
            "timed_out": timed_out,
            "stdout": stdout[:MAX_OUTPUT_BYTES].decode("utf-8", errors="replace"),
            "stderr": stderr[:MAX_OUTPUT_BYTES].decode("utf-8", errors="replace"),
            "output_truncated": (
                len(stdout) > MAX_OUTPUT_BYTES or len(stderr) > MAX_OUTPUT_BYTES
            ),
        }
    finally:
        try:
            os.unlink(script_path)
        except FileNotFoundError:
            pass


class Handler(BaseHTTPRequestHandler):
    """Serve health and bounded Python execution requests."""

    server_version = "HermesChatSandbox/1"

    def log_message(self, format_string: str, *args: object) -> None:
        """Keep normal request logs concise and free of request bodies."""
        print(f"sandbox: {self.address_string()} {format_string % args}", flush=True)

    def _json(self, status: int, payload: dict[str, object]) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("X-Content-Type-Options", "nosniff")
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self) -> None:
        """Return a minimal unauthenticated health response."""
        if self.path == "/health":
            self._json(200, {"status": "ok"})
        else:
            self._json(404, {"error": "not found"})

    def do_POST(self) -> None:
        """Validate and execute a Python request from the matching tenant pod."""
        if self.path != "/v1/execute":
            self._json(404, {"error": "not found"})
            return
        try:
            length = int(self.headers.get("Content-Length", "0"))
        except ValueError:
            self._json(400, {"error": "invalid content length"})
            return
        if length <= 0 or length > MAX_REQUEST_BYTES:
            self._json(413, {"error": "request too large"})
            return
        try:
            payload = json.loads(self.rfile.read(length))
        except (json.JSONDecodeError, UnicodeDecodeError):
            self._json(400, {"error": "invalid JSON"})
            return
        code = payload.get("code") if isinstance(payload, dict) else None
        if not isinstance(code, str) or not code.strip():
            self._json(400, {"error": "code is required"})
            return
        if len(code.encode("utf-8")) > MAX_CODE_BYTES:
            self._json(413, {"error": "code exceeds 128 KiB"})
            return
        self._json(200, _execute(code))


if __name__ == "__main__":
    WORKSPACE.mkdir(parents=True, exist_ok=True)
    ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()