"""Linux code execution isolated with Bubblewrap namespaces."""

from __future__ import annotations

import asyncio
import os
import shutil
import tempfile
import time
from pathlib import Path
from typing import Any, Final

from cyberbrain26.product.local_store import (
    ENV_WORKSPACES_ROOT,
    ensure_workspaces_root,
    resolve_workspace,
)

# The host resolver file, bound read-only into the sandbox only when an
# authenticated engagement profile authorised egress. A shared network
# namespace brings routes but no /etc, so without this one file a target named
# in the engagement could not be resolved at all.
RESOLVER_PATH: Final[str] = "/etc/resolv.conf"
# Ceiling for the /tmp tmpfs, applied only in that same profile. The default
# tmpfs is half of RAM, which on a large host is unbounded from the engagement's
# point of view; scan tools write their findings to /tmp faster than the
# wall-clock ceiling expires.
NETWORK_TMPFS_SIZE_BYTES: Final[int] = 512 * 1024 * 1024
# Process ceiling for that profile. Sharing the host's network namespace is the
# one place where forking in parallel becomes possible, so the fork-bomb
# ceiling that keeps a sandboxed process from exhausting the host's PID
# namespace is pinned here as well.
NETWORK_MAX_PROCESSES: Final[int] = 256


class BubblewrapExecutionProvider:
    """Run Python or shell code in an isolated, disposable Linux filesystem.

    The sandbox gets read-only access to system runtime directories, a private
    temporary directory, and one writable workspace mount. Network, process,
    IPC, and other namespaces are unshared. There is deliberately no
    unsandboxed fallback when Bubblewrap is unavailable or cannot start.

    ``allow_network`` is the one profile that changes that, and it changes
    exactly three things: the network namespace is *shared* with the host instead
    of unshared, the ``/etc/resolv.conf`` of the host is bound read-only so
    a target named in the engagement can be resolved at all, and the ``/tmp``
    tmpfs gets a size ceiling (with the matching ``RLIMIT_NPROC``). Nothing
    else about the sandbox differs. Every filesystem, PID, IPC and UTS
    isolation, the read-only ``/usr``, the disposable ``/tmp`` and the single
    writable workspace bind stay exactly as they are, and the flag is ``False``
    by default, so a caller that does not ask for it gets byte-for-byte the
    argument vector this class has always produced. It is set only from an
    authenticated engagement profile (see
    ``execution_plane.secure_gateway.engagement_profile_allows_network``) and
    only for the pentest role; it is not a general "let code use the internet"
    switch, and it says nothing about untrusted code - a process in a shared
    network namespace is still confined by everything else listed here, and
    still must not be described as safe to run hostile code.

    The namespace part of that claim is only true because the vector is
    ``--unshare-all --share-net`` and not ``--share-net``. bwrap applies no
    unsharing by default, so omitting ``--unshare-all`` silently shares the
    host's PID namespace too - which is a credential-exfiltration primitive, not
    a weaker sandbox. The regression test for this is
    ``test_the_network_profile_keeps_its_own_pid_namespace`` in
    ``tests/test_execution_provider.py``, which launches bwrap and reads
    ``/proc/1/comm``; an argv-string assertion cannot catch this class of bug.

    The writable workspace is **not** an arbitrary host directory. Because the
    workspace is bind-mounted read-write, a workspace such as ``/root/.ssh``
    would undo every other control in this class: read-only ``/usr``, disposable
    ``/tmp`` and unshared namespaces all keep their value only if the code
    cannot reach the rest of the host through the one mount it is given. A
    workspace is therefore confined to the sanctioned root
    ``<data dir>/workspaces`` (:func:`~cyberbrain26.product.local_store.workspaces_root`).
    The default root is the single source of truth for the data directory, so
    ``CYBERBRAIN_DATA_DIR`` moves it. Deployments that genuinely need another
    root pass ``workspaces_root`` explicitly or set
    ``CYBERBRAIN_WORKSPACES_ROOT``; widening the confinement is then the
    caller's explicit decision, recorded in configuration, never a silent
    fallback to the old permissive behaviour.
    """

    def __init__(
        self,
        *,
        bubblewrap_binary: str = "bwrap",
        python_executable: str = "/usr/bin/python3",
        shell_executable: str = "/usr/bin/sh",
        max_timeout_seconds: float = 30.0,
        max_memory_mb: int = 512,
        max_cpu_seconds: int = 10,
        max_file_size_mb: int = 16,
        max_output_bytes: int = 1_000_000,
        workspaces_root: str | Path | None = None,
        allow_network: bool = False,
    ) -> None:
        if max_timeout_seconds <= 0:
            raise ValueError("max_timeout_seconds must be positive")
        if max_memory_mb <= 0 or max_cpu_seconds <= 0 or max_file_size_mb <= 0:
            raise ValueError("resource limits must be positive")
        if max_output_bytes <= 0:
            raise ValueError("max_output_bytes must be positive")
        self.bubblewrap_binary = bubblewrap_binary
        self.python_executable = python_executable
        self.shell_executable = shell_executable
        self.max_timeout_seconds = max_timeout_seconds
        self.max_memory_mb = max_memory_mb
        self.max_cpu_seconds = max_cpu_seconds
        self.max_file_size_mb = max_file_size_mb
        self.max_output_bytes = max_output_bytes
        self._workspaces_root = Path(workspaces_root).expanduser() if workspaces_root else None
        self.allow_network = allow_network

    @property
    def provider_name(self) -> str:
        return "bubblewrap"

    @property
    def allowed_workspaces_root(self) -> Path:
        """Return the directory a writable workspace must live inside.

        Precedence: the ``workspaces_root`` constructor argument, then the
        ``CYBERBRAIN_WORKSPACES_ROOT`` environment variable, then the default
        sanctioned root under the data directory. The default root is created
        on demand with owner-only permissions so a fresh installation works;
        an explicitly configured root is only resolved, never created, because
        it belongs to the caller.

        The value is resolved per call rather than at construction time so the
        root always reflects the current environment.
        """
        if self._workspaces_root is not None:
            return self._workspaces_root.resolve()
        env_value = os.environ.get(ENV_WORKSPACES_ROOT)
        if env_value is not None and env_value.strip():
            return Path(env_value.strip()).expanduser().resolve()
        return ensure_workspaces_root()

    async def execute(
        self,
        code: str,
        language: str,
        *,
        timeout_seconds: int = 30,
        workspace: str | None = None,
        allow_network: bool = False,
    ) -> dict[str, Any]:
        """Execute code and return captured output and process status.

        ``allow_network`` mirrors the constructor argument for the one caller
        that decides per call rather than per provider
        (``AuthorizedExecutionService``, which derives it from the
        authenticated engagement profile). The two are combined with ``or``, so
        a call can grant what the constructor granted but never take a grant
        away: the narrower of the two is always the default, and neither can
        widen the sandbox on its own.
        """
        network_allowed = self.allow_network or allow_network
        language_key = language.strip().lower()
        if language_key in {"python", "py"}:
            command = [self.python_executable, "-I", "-"]
            stdin = code.encode("utf-8")
        elif language_key in {"shell", "sh"}:
            command = [self.shell_executable, "-c", code]
            stdin = None
        else:
            raise ValueError(f"unsupported execution language: {language}")

        # The workspace comes from a step payload, so it is validated before
        # anything is launched: an unconfined workspace must never reach the
        # point where a process exists, let alone one with a read-write bind
        # mount of the host.
        workspace_path: Path | None = None
        if workspace is not None:
            workspace_path = Path(workspace).expanduser().resolve(strict=True)
            workspace_path = resolve_workspace(
                workspace_path, root=self.allowed_workspaces_root
            )
            if not workspace_path.is_dir():
                raise ValueError("workspace must be an existing directory")

        binary = shutil.which(self.bubblewrap_binary)
        if binary is None:
            raise RuntimeError(
                f"Bubblewrap executable {self.bubblewrap_binary!r} is unavailable; "
                "refusing to run code without a sandbox"
            )
        timeout = min(float(timeout_seconds), self.max_timeout_seconds)
        if timeout <= 0:
            raise ValueError("timeout_seconds must be positive")

        if workspace_path is None:
            with tempfile.TemporaryDirectory(prefix="cyberbrain26-exec-") as temp_dir:
                return await self._run(
                    binary, command, stdin, Path(temp_dir), timeout,
                    allow_network=network_allowed,
                )
        return await self._run(
            binary, command, stdin, workspace_path, timeout,
            allow_network=network_allowed,
        )

    async def _run(
        self,
        binary: str,
        command: list[str],
        stdin: bytes | None,
        workspace: Path,
        timeout: float,
        *,
        allow_network: bool = False,
    ) -> dict[str, Any]:
        bwrap_args = [
            binary,
            # bwrap takes ``--unshare-all`` first and then ``--share-net`` to
            # re-share ONLY the network namespace, overriding the earlier flag.
            # Passing ``--share-net`` on its own - which is what this did until
            # 2026-10-03 - drops EVERY unshare directive, because bwrap has no
            # default for them: the sandbox then shared the host's PID, IPC, UTS
            # and user namespaces. Verified by running bwrap, not by reading it:
            #   --share-net        -> /proc/1/comm == "systemd"  (host init)
            #   --unshare-all      -> /proc/1/comm == "bwrap"   (sandbox init)
            # and with the former a scan could list 327 host processes and read
            # /proc/<pid>/environ of a host process - i.e. every API key and
            # token belonging to this user - straight out of the scan sandbox.
            # With ``allow_network`` False - the default, and what every role
            # other than an authenticated pentest gets - this is the exact
            # argument list this method has always produced.
            *(
                ("--unshare-all", "--share-net")
                if allow_network
                else ("--unshare-all",)
            ),
            "--die-with-parent",
            "--new-session",
            "--ro-bind", "/usr", "/usr",
        ]
        for runtime_path in ("/lib", "/lib64"):
            if Path(runtime_path).exists():
                bwrap_args.extend(("--ro-bind", runtime_path, runtime_path))
        if allow_network and Path(RESOLVER_PATH).is_file():
            # A shared network namespace brings the host's routes but not the
            # host's resolver, and the sandbox has no /etc at all, so without
            # this one read-only file a target the engagement authorised *by
            # name* could not be resolved and every such scan would fail with a
            # name-resolution error rather than a refusal. It is the only piece
            # of host configuration bound in, and it is bound read-only. When
            # the host has no resolver file at all the bind is skipped rather
            # than guessed at: IP-literal targets still work, and a name simply
            # does not resolve, which is a visible failure and not a wider grant.
            bwrap_args.extend(("--ro-bind", RESOLVER_PATH, RESOLVER_PATH))
        # bwrap parses ``--size`` as a modifier of the tmpfs that FOLLOWS it, so
        # the ceiling has to be emitted before the tmpfs it applies to:
        # ``--size N --tmpfs /tmp``. The reverse order is not a silently ignored
        # flag - bwrap refuses it and exits ("--size must be followed by
        # --tmpfs"), which is how this profile shipped broken with no test to
        # notice. With ``allow_network`` False - the default, and what every role
        # other than an authenticated pentest gets - this stays exactly
        # ``["--tmpfs", "/tmp"]`` at exactly this position, so the default
        # vector stays byte-for-byte what this class has always produced.
        tmpfs_args: list[str] = ["--tmpfs", "/tmp"]
        if allow_network:
            # Scan tools write their findings to /tmp faster than the wall-clock
            # ceiling expires; the default tmpfs is half of RAM, which on a
            # large host is unbounded from the engagement's point of view. The
            # ceiling applies to the tmpfs emitted here and only in this profile.
            tmpfs_args = ["--size", str(NETWORK_TMPFS_SIZE_BYTES), "--tmpfs", "/tmp"]
        bwrap_args.extend(
            (
                "--dev", "/dev",
                "--proc", "/proc",
                *tmpfs_args,
                "--bind", os.fspath(workspace), "/workspace",
                "--chdir", "/workspace",
                "--clearenv",
                "--setenv", "PATH", "/usr/bin:/bin",
                "--setenv", "HOME", "/tmp",
                "--",
                *self._limited_command(command, allow_network=allow_network),
            )
        )

        started = time.monotonic()
        process = await asyncio.create_subprocess_exec(
            *bwrap_args,
            stdin=asyncio.subprocess.PIPE if stdin is not None else asyncio.subprocess.DEVNULL,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        assert process.stdout is not None
        assert process.stderr is not None
        stdout_task = asyncio.create_task(self._read_capped(process.stdout))
        stderr_task = asyncio.create_task(self._read_capped(process.stderr))
        timed_out = False
        try:
            if stdin is not None:
                assert process.stdin is not None
                process.stdin.write(stdin)
                await process.stdin.drain()
                process.stdin.close()
            await asyncio.wait_for(process.wait(), timeout=timeout)
        except TimeoutError:
            timed_out = True
            process.kill()
            await process.wait()
        except asyncio.CancelledError:
            process.kill()
            await process.wait()
            await asyncio.gather(stdout_task, stderr_task, return_exceptions=True)
            raise

        (stdout, stdout_truncated), (stderr, stderr_truncated) = await asyncio.gather(
            stdout_task, stderr_task
        )
        return {
            "provider": self.provider_name,
            "stdout": stdout.decode("utf-8", errors="replace"),
            "stderr": stderr.decode("utf-8", errors="replace"),
            "return_code": process.returncode,
            "timed_out": timed_out,
            "stdout_truncated": stdout_truncated,
            "stderr_truncated": stderr_truncated,
            "duration_seconds": time.monotonic() - started,
        }

    def _limited_command(
        self, command: list[str], *, allow_network: bool = False
    ) -> list[str]:
        """Apply CPU, address-space, and output-file limits before exec.

        With ``allow_network`` False - the default - the source string below is
        byte-for-byte the one this method has always passed to ``-c``. The only
        difference in the other profile is the extra ``RLIMIT_NPROC`` block
        spliced into that same source text, never an extra argv element: an
        extra slot would shift ``sys.argv[4]`` and make the limiter exec itself
        instead of the requested program.
        """
        limiter = (
            "import os, resource, sys\n"
            "cpu, memory, file_size = map(int, sys.argv[1:4])\n"
            "for key, limit in ((resource.RLIMIT_CPU, cpu), "
            "(resource.RLIMIT_AS, memory), (resource.RLIMIT_FSIZE, file_size)):\n"
            "    _, hard = resource.getrlimit(key)\n"
            "    if hard != resource.RLIM_INFINITY:\n"
            "        limit = min(limit, hard)\n"
            "    resource.setrlimit(key, (limit, limit))\n"
        )
        if allow_network:
            limiter += (
                f"nproc = {NETWORK_MAX_PROCESSES}\n"
                "_, hard = resource.getrlimit(resource.RLIMIT_NPROC)\n"
                "if hard != resource.RLIM_INFINITY:\n"
                "    nproc = min(nproc, hard)\n"
                "try:\n"
                "    resource.setrlimit(resource.RLIMIT_NPROC, (nproc, nproc))\n"
                "except (OSError, ValueError) as exc:\n"
                "    sys.stderr.write('RLIMIT_NPROC: ' + str(exc) + '\\n')\n"
                "    sys.stderr.flush()\n"
                "    os._exit(126)\n"
            )
        limiter += (
            "executable = sys.argv[4]\n"
            "os.execv(executable, [executable, *sys.argv[5:]])\n"
        )
        return [
            self.python_executable,
            "-I",
            "-c",
            limiter,
            str(self.max_cpu_seconds),
            str(self.max_memory_mb * 1024 * 1024),
            str(self.max_file_size_mb * 1024 * 1024),
            command[0],
            *command[1:],
        ]

    async def _read_capped(self, stream: asyncio.StreamReader) -> tuple[bytes, bool]:
        chunks: list[bytes] = []
        retained = 0
        truncated = False
        while chunk := await stream.read(65536):
            remaining = self.max_output_bytes - retained
            if remaining > 0:
                kept = chunk[:remaining]
                chunks.append(kept)
                retained += len(kept)
            if len(chunk) > max(remaining, 0):
                truncated = True
        return b"".join(chunks), truncated


__all__ = ["BubblewrapExecutionProvider"]