core/docker_executor.py

"""Docker-based sandbox executor for ODSE.

Executes Python code in isolated Docker containers with strict security constraints:
- Resource limits (CPU, memory)
- Read-only root filesystem (except for specific mounts)
- No network access
- No privilege escalation
- Whitelisted imports enforced
- Timeout enforcement via container timeout
"""

from __future__ import annotations

import pickle
import shutil
import subprocess
import tempfile
import time
import traceback
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

import numpy as np
import pandas as pd

from .executor import ExecutionResult, ALLOWED_MODULES
from .models import ExecutionStatus, VariableInfo


class DockerSandboxExecutor:
    """Executes Python code in isolated Docker containers."""

    def __init__(
        self,
        timeout_seconds: float = 30.0,
        max_output_chars: int = 10_000,
        docker_image: str = "odse-sandbox:latest",
        memory_limit: str = "512m",
        cpus: float = 1.0,
    ) -> None:
        self.timeout_seconds = timeout_seconds
        self.max_output_chars = max_output_chars
        self.docker_image = docker_image
        self.memory_limit = memory_limit
        self.cpus = cpus
        self._namespace: Dict[str, Any] = {}
        self._setup_done: bool = False
        self._work_dir: Optional[Path] = None
        self._namespace_file: Optional[Path] = None
        self._evaluate_fn: Optional[Callable] = None

        # Validate Docker is available
        try:
            subprocess.run(
                ["docker", "info"],
                capture_output=True,
                timeout=5,
                check=True,
            )
        except Exception as e:
            raise RuntimeError(
                f"Docker is not available or not running: {e}. "
                "Ensure Docker is installed and the daemon is running."
            ) from e

    @property
    def namespace(self) -> Dict[str, Any]:
        """Direct (read-only) view of the sandbox namespace."""
        return self._namespace

    def setup_namespace(
        self,
        *,
        train_df: pd.DataFrame,
        val_features: pd.DataFrame,
        test_features: pd.DataFrame,
        target_column: str,
        evaluate_fn: Callable,
    ) -> None:
        """Initialise the sandbox namespace with pre-loaded variables."""
        self._work_dir = Path(tempfile.mkdtemp(prefix="odse_sandbox_"))
        self._evaluate_fn = evaluate_fn

        self._namespace = {
            "train_df": train_df.copy(),
            "val_features": val_features.copy(),
            "test_features": test_features.copy(),
            "target_column": target_column,
            "pd": pd,
            "np": np,
            "evaluate": None,  # Will be injected as code
            "print": print,
        }

        self._namespace_file = self._work_dir / "namespace.pkl"
        self._save_namespace()
        self._setup_done = True

    def reset(self) -> None:
        """Clear the namespace and clean up temporary files."""
        self._namespace.clear()
        self._setup_done = False
        if self._work_dir and self._work_dir.exists():
            shutil.rmtree(self._work_dir, ignore_errors=True)
        self._work_dir = None
        self._namespace_file = None

    def execute(self, code: str) -> ExecutionResult:
        """Execute *code* in a Docker container and return an ``ExecutionResult``."""
        if not self._setup_done:
            return ExecutionResult(
                status=ExecutionStatus.ERROR,
                stderr="Sandbox not initialised - call setup_namespace() first.",
            )

        start = time.perf_counter()

        try:
            self._save_namespace()

            code_file = self._work_dir / "code.py"
            evaluate_setup_code = self._get_evaluate_setup_code()
            full_code = evaluate_setup_code + "\n\n" + code
            code_file.write_text(full_code, encoding="utf-8")

            cmd = [
                "docker", "run",
                "--rm",
                f"--memory={self.memory_limit}",
                f"--cpus={self.cpus}",
                "--network=none",
                "--read-only",
                f"--tmpfs=/tmp:size=100m,mode=1777",
                f"--tmpfs=/home:size=100m,mode=1777",
                "-v", f"{self._work_dir}:/sandbox/work:rw",
                "--cap-drop=ALL",
                "--security-opt=no-new-privileges",
                self.docker_image,
                "timeout", str(int(self.timeout_seconds) + 5),
                "python", "/app/sandbox_runner.py",
                "/sandbox/work/namespace.pkl",
                "/sandbox/work/code.py",
            ]

            result = subprocess.run(
                cmd,
                capture_output=True,
                timeout=int(self.timeout_seconds) + 10,
                text=True,
            )

            elapsed = (time.perf_counter() - start) * 1000

            if result.returncode == 124:
                return ExecutionResult(
                    status=ExecutionStatus.TIMEOUT,
                    stdout=self._truncate(result.stdout),
                    stderr=f"Code execution exceeded {self.timeout_seconds}s time limit",
                    execution_time_ms=elapsed,
                )

            if self._namespace_file and self._namespace_file.exists():
                self._load_namespace()

            if result.returncode == 0:
                return ExecutionResult(
                    status=ExecutionStatus.SUCCESS,
                    stdout=self._truncate(result.stdout),
                    stderr="",
                    execution_time_ms=elapsed,
                )
            else:
                return ExecutionResult(
                    status=ExecutionStatus.ERROR,
                    stdout="",
                    stderr=self._truncate(result.stderr or result.stdout),
                    execution_time_ms=elapsed,
                )

        except subprocess.TimeoutExpired:
            elapsed = (time.perf_counter() - start) * 1000
            return ExecutionResult(
                status=ExecutionStatus.TIMEOUT,
                stderr="Container execution exceeded timeout",
                execution_time_ms=elapsed,
            )

        except Exception:
            elapsed = (time.perf_counter() - start) * 1000
            return ExecutionResult(
                status=ExecutionStatus.ERROR,
                stderr=self._truncate(traceback.format_exc()),
                execution_time_ms=elapsed,
            )

    def get_namespace_summary(self) -> List[VariableInfo]:
        """Return a summary of user-visible variables in the namespace."""
        hidden = {"__builtins__", "pd", "np", "evaluate", "target_column", "print"}
        summary: List[VariableInfo] = []
        for name, value in self._namespace.items():
            if name.startswith("_") or name in hidden:
                continue
            summary.append(
                VariableInfo(
                    name=name,
                    type_name=type(value).__name__,
                    shape=getattr(value, "shape", None),
                    preview=self._preview(value),
                )
            )
        return summary

    def get_predictions(self) -> Optional[np.ndarray]:
        """Retrieve ``predictions`` from the namespace (or ``None``)."""
        preds = self._namespace.get("predictions")
        if preds is None:
            return None
        try:
            return np.asarray(preds)
        except Exception:
            return None

    def __del__(self) -> None:
        """Cleanup on deletion."""
        self.reset()

    # -- Private helpers ----

    def _get_evaluate_setup_code(self) -> str:
        """Generate code to set up the evaluate function in the container."""
        evaluate_fn = self._evaluate_fn
        
        if not (hasattr(evaluate_fn, 'val_labels') and hasattr(evaluate_fn, 'problem_type') and hasattr(evaluate_fn, 'metric')):
            return ""

        problem_type = evaluate_fn.problem_type.value
        metric = evaluate_fn.metric

        # Use string concatenation to avoid nested f-string issues
        # Note: _val_labels will be pre-loaded in the namespace by _save_namespace()
        setup_code = (
            "import numpy as np\n"
            "import sklearn.metrics\n\n"
            "def evaluate(predictions):\n"
            "    preds = np.asarray(predictions)\n"
            "    if len(preds) != len(_val_labels):\n"
            "        return {'error': f'Expected {len(_val_labels)} predictions (val_features length), got {len(preds)}'}\n"
            "    \n"
            "    try:\n"
            + f"        if '{problem_type}' == 'classification':\n"
            "            acc = sklearn.metrics.accuracy_score(_val_labels, preds)\n"
            "            primary = acc\n"
            "            report = {'accuracy': round(acc, 4), 'f1_macro': round(sklearn.metrics.f1_score(_val_labels, preds, average='macro', zero_division=0), 4)}\n"
            "        else:\n"
            "            r2 = sklearn.metrics.r2_score(_val_labels, preds)\n"
            "            primary = r2\n"
            "            report = {'r2': round(r2, 4)}\n"
            "        \n"
            + f"        report['primary_metric'] = '{metric}'\n"
            "        report['primary_score'] = round(primary, 4)\n"
            "        return report\n"
            "    except Exception as e:\n"
            "        return {'error': str(e)}\n"
        )
        return setup_code

    def _save_namespace(self) -> None:
        """Serialize the current namespace to a pickle file."""
        if not self._namespace_file:
            raise RuntimeError("Namespace file not initialized")

        pickleable = {}
        for key, value in self._namespace.items():
            if key in ("print", "evaluate"):
                continue
            try:
                pickle.dumps(value)
                pickleable[key] = value
            except (pickle.PicklingError, TypeError):
                pass

        # Add validation labels to the namespace so the evaluate function can use it
        if self._evaluate_fn and hasattr(self._evaluate_fn, 'val_labels'):
            try:
                pickle.dumps(self._evaluate_fn.val_labels)
                pickleable['_val_labels'] = self._evaluate_fn.val_labels
            except (pickle.PicklingError, TypeError):
                pass

        with open(self._namespace_file, "wb") as f:
            pickle.dump(pickleable, f)

    def _load_namespace(self) -> None:
        """Load the namespace from the pickle file written by the container."""
        if not self._namespace_file or not self._namespace_file.exists():
            return

        try:
            with open(self._namespace_file, "rb") as f:
                loaded = pickle.load(f)
                print_fn = self._namespace.get("print")
                evaluate_fn = self._evaluate_fn
                self._namespace.update(loaded)
                if print_fn:
                    self._namespace["print"] = print_fn
                if evaluate_fn:
                    self._namespace["evaluate"] = evaluate_fn
        except Exception:
            pass

    def _truncate(self, text: str) -> str:
        if len(text) <= self.max_output_chars:
            return text
        return text[: self.max_output_chars] + "\n... [output truncated]"

    @staticmethod
    def _preview(value: Any, max_len: int = 300) -> str:
        """Generate a short string preview of *value*."""
        try:
            if isinstance(value, pd.DataFrame):
                return f"DataFrame(shape={value.shape}, cols={list(value.columns[:5])})"
            if isinstance(value, pd.Series):
                return f"Series(len={len(value)}, dtype={value.dtype})"
            if isinstance(value, np.ndarray):
                return f"ndarray(shape={value.shape}, dtype={value.dtype})"
            s = repr(value)
            return s[:max_len] if len(s) > max_len else s
        except Exception:
            return "<unprintable>"