container_shell: isolated network by default, plus optional network transport

src/seclab_taskflows/mcp_servers/container_shell.py

@@ -1,6 +1,45 @@
 # SPDX-FileCopyrightText: GitHub, Inc.
 # SPDX-License-Identifier: MIT
 
+"""MCP server that runs shell commands inside a managed Docker container.
+
+Configuration is read from the process environment (set per toolbox in the
+toolbox YAML's ``server_params.env`` block):
+
+- ``CONTAINER_IMAGE`` — image to run (required).
+- ``CONTAINER_WORKSPACE`` — host path bind-mounted at ``/workspace`` (optional).
+- ``CONTAINER_TIMEOUT`` — default per-command timeout in seconds (default 30).
+- ``CONTAINER_PERSIST`` — reuse a deterministic container across runs when truthy.
+- ``CONTAINER_PERSIST_KEY`` — extra key to distinguish persistent containers.
+- ``CONTAINER_NETWORK`` — Docker network mode for the container. Defaults to
+  ``none`` so the container is egress-locked. Set it to ``bridge``, ``host``, or
+  a user-defined network to enable networking.
+
+Selecting a network mode from a toolbox: the agent passes only the toolbox's
+declared ``env`` entries to this server, so a network mode is selectable at run
+time only if the toolbox exposes the knob. To let callers opt in, add a
+passthrough line to the toolbox ``env`` block, e.g.::
+
+    CONTAINER_NETWORK: "{{ env('CONTAINER_NETWORK', 'none') }}"
+
+A toolbox that needs networking by default (e.g. recon tooling) can use
+``'bridge'`` as the template default instead. An empty or unset value always
+falls back to ``none``, so isolation cannot be disabled by a blank variable.
+
+Transport:
+
+- ``CONTAINER_SHELL_TRANSPORT`` — MCP transport. Defaults to ``stdio`` for the
+  standard local case where the agent launches this server as a subprocess. Set
+  it to ``http``, ``streamable-http``, or ``sse`` to run as a network-accessible
+  MCP server that a remote agent can connect to.
+- ``CONTAINER_SHELL_HOST`` / ``CONTAINER_SHELL_PORT`` — bind address for the
+  network transports (defaults ``127.0.0.1`` / ``8080``); ignored for ``stdio``.
+
+Which Docker daemon this server drives is orthogonal to the transport: the
+docker CLI honours ``DOCKER_HOST`` from the environment, so pointing this server
+at a specific (e.g. dedicated, isolated) daemon needs no code change here.
+"""
+
 import atexit
 import hashlib
 import json
@@ -30,6 +69,22 @@
 CONTAINER_TIMEOUT = int(os.environ.get("CONTAINER_TIMEOUT", "30"))
 CONTAINER_PERSIST = os.environ.get("CONTAINER_PERSIST", "").lower() in ("1", "true", "yes")
 CONTAINER_PERSIST_KEY = os.environ.get("CONTAINER_PERSIST_KEY", "")
+# Docker network mode for the container. Defaults to "none" so containers are
+# egress-locked (no network access) unless a caller explicitly opts in by
+# setting CONTAINER_NETWORK to a network name such as "bridge", "host", or a
+# user-defined network. An empty or whitespace-only value falls back to "none"
+# so the isolation default cannot be silently disabled by an unset variable.
+CONTAINER_NETWORK = os.environ.get("CONTAINER_NETWORK", "none").strip() or "none"
+# MCP transport selection. Defaults to "stdio" for the standard local case
+# where the agent launches this server as a subprocess. Set
+# CONTAINER_SHELL_TRANSPORT to "http", "streamable-http", or "sse" to run as a
+# network-accessible server (for example, an isolated sidecar reached by a
+# remote agent); CONTAINER_SHELL_HOST/PORT control the bind address for those
+# transports and are ignored for stdio.
+CONTAINER_SHELL_TRANSPORT = os.environ.get("CONTAINER_SHELL_TRANSPORT", "stdio").strip() or "stdio"
+CONTAINER_SHELL_HOST = os.environ.get("CONTAINER_SHELL_HOST", "127.0.0.1")
+CONTAINER_SHELL_PORT = int(os.environ.get("CONTAINER_SHELL_PORT", "8080"))
+_SUPPORTED_TRANSPORTS = ("stdio", "http", "streamable-http", "sse")
 
 _DEFAULT_WORKDIR = "/workspace"
 _DOCKER_TIMEOUT = 30
@@ -38,11 +93,15 @@
 def _persistent_name() -> str:
     """Derive a deterministic container name from the image for reuse across tasks.
 
-    Incorporates a hash of the full image reference (and optional
-    CONTAINER_PERSIST_KEY) to avoid collisions between long image names that
-    share a common prefix, or between independent runs of the same image.
+    Incorporates a hash of the full image reference, the configured network
+    mode, and an optional CONTAINER_PERSIST_KEY. Including the network mode
+    ensures a run configured for one network (e.g. the default "none") never
+    reuses a persistent container that was created with a different, more
+    permissive network (e.g. "bridge"), which would otherwise silently
+    re-enable egress. The hash also avoids collisions between long image names
+    that share a common prefix, or between independent runs of the same image.
     """
-    key_material = CONTAINER_IMAGE
+    key_material = f"{CONTAINER_IMAGE}:net={CONTAINER_NETWORK}"
     if CONTAINER_PERSIST_KEY:
         key_material += f":{CONTAINER_PERSIST_KEY}"
     digest = hashlib.sha256(key_material.encode()).hexdigest()[:12]
@@ -106,7 +165,7 @@ def _start_container() -> str:
     else:
         name = f"seclab-shell-{uuid.uuid4().hex[:8]}"
 
-    cmd = ["docker", "run", "-d", "--name", name]
+    cmd = ["docker", "run", "-d", "--name", name, "--network", CONTAINER_NETWORK]
     if not CONTAINER_PERSIST:
         cmd.append("--rm")
     if CONTAINER_WORKSPACE:
@@ -176,5 +235,29 @@ def shell_exec(
     return output
 
 
+def _run_server() -> None:
+    """Run the MCP server using the configured transport.
+
+    Defaults to stdio (local subprocess use). When CONTAINER_SHELL_TRANSPORT
+    selects a network transport, the server binds CONTAINER_SHELL_HOST:PORT so a
+    remote agent can reach it.
+    """
+    if CONTAINER_SHELL_TRANSPORT not in _SUPPORTED_TRANSPORTS:
+        msg = (
+            f"Unsupported CONTAINER_SHELL_TRANSPORT {CONTAINER_SHELL_TRANSPORT!r}; "
+            f"expected one of {', '.join(_SUPPORTED_TRANSPORTS)}"
+        )
+        raise ValueError(msg)
+    if CONTAINER_SHELL_TRANSPORT == "stdio":
+        mcp.run(show_banner=False)
+    else:
+        mcp.run(
+            transport=CONTAINER_SHELL_TRANSPORT,
+            host=CONTAINER_SHELL_HOST,
+            port=CONTAINER_SHELL_PORT,
+            show_banner=False,
+        )
+
+
 if __name__ == "__main__":
-    mcp.run(show_banner=False)
+    _run_server()

src/seclab_taskflows/toolboxes/container_shell_base.yaml

@@ -15,6 +15,7 @@ server_params:
     CONTAINER_TIMEOUT: "{{ env('CONTAINER_TIMEOUT', '30') }}"
     CONTAINER_PERSIST: "{{ env('CONTAINER_PERSIST', required=False) }}"
     CONTAINER_PERSIST_KEY: "{{ env('CONTAINER_PERSIST_KEY', required=False) }}"
+    CONTAINER_NETWORK: "{{ env('CONTAINER_NETWORK', required=False) }}"
     LOG_DIR: "{{ env('LOG_DIR') }}"
 
 confirm:

src/seclab_taskflows/toolboxes/container_shell_malware_analysis.yaml

@@ -15,6 +15,7 @@ server_params:
     CONTAINER_TIMEOUT: "{{ env('CONTAINER_TIMEOUT', '60') }}"
     CONTAINER_PERSIST: "{{ env('CONTAINER_PERSIST', required=False) }}"
     CONTAINER_PERSIST_KEY: "{{ env('CONTAINER_PERSIST_KEY', required=False) }}"
+    CONTAINER_NETWORK: "{{ env('CONTAINER_NETWORK', required=False) }}"
     LOG_DIR: "{{ env('LOG_DIR') }}"
 
 confirm:

src/seclab_taskflows/toolboxes/container_shell_network_analysis.yaml

@@ -15,6 +15,7 @@ server_params:
     CONTAINER_TIMEOUT: "{{ env('CONTAINER_TIMEOUT', '30') }}"
     CONTAINER_PERSIST: "{{ env('CONTAINER_PERSIST', required=False) }}"
     CONTAINER_PERSIST_KEY: "{{ env('CONTAINER_PERSIST_KEY', required=False) }}"
+    CONTAINER_NETWORK: "{{ env('CONTAINER_NETWORK', 'bridge') }}"
     LOG_DIR: "{{ env('LOG_DIR') }}"
 
 confirm:

src/seclab_taskflows/toolboxes/container_shell_sast.yaml

@@ -15,6 +15,7 @@ server_params:
     CONTAINER_TIMEOUT: "{{ env('CONTAINER_TIMEOUT', '60') }}"
     CONTAINER_PERSIST: "{{ env('CONTAINER_PERSIST', required=False) }}"
     CONTAINER_PERSIST_KEY: "{{ env('CONTAINER_PERSIST_KEY', required=False) }}"
+    CONTAINER_NETWORK: "{{ env('CONTAINER_NETWORK', required=False) }}"
     LOG_DIR: "{{ env('LOG_DIR') }}"
 
 confirm:

tests/test_container_shell.py

@@ -2,6 +2,9 @@
 # SPDX-License-Identifier: MIT
 
 import subprocess
+import importlib
+import os
+import atexit
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -28,6 +31,30 @@ def _reset_container():
     cs_mod._container_name = None
 
 
+def _reload_cs():
+    """Reload cs_mod without accumulating atexit handlers.
+
+    The module registers ``_stop_container`` with ``atexit`` at import time, so
+    reloading would stack a fresh handler each time. Unregister the current one
+    first so exactly one handler remains after the reload.
+    """
+    atexit.unregister(cs_mod._stop_container)
+    return importlib.reload(cs_mod)
+
+
+def _restore_env_and_reload(var, original):
+    """Restore an env var to its pre-test value and reload cs_mod.
+
+    Reloading with the real environment (rather than the monkeypatched value)
+    keeps the module-level config from leaking into subsequent tests.
+    """
+    if original is None:
+        os.environ.pop(var, None)
+    else:
+        os.environ[var] = original
+    _reload_cs()
+
+
 # ---------------------------------------------------------------------------
 # _start_container tests
 # ---------------------------------------------------------------------------
@@ -89,6 +116,48 @@ def test_start_container_rejects_empty_image(self):
             with pytest.raises(RuntimeError, match="CONTAINER_IMAGE is not set"):
                 cs_mod._start_container()
 
+    def test_start_container_default_network_none(self):
+        with (
+            patch.object(cs_mod, "CONTAINER_IMAGE", "test-image:latest"),
+            patch.object(cs_mod, "CONTAINER_WORKSPACE", ""),
+            patch.object(cs_mod, "CONTAINER_NETWORK", "none"),
+            patch("subprocess.run", return_value=_make_proc(returncode=0)) as mock_run,
+        ):
+            cs_mod._start_container()
+            cmd = mock_run.call_args[0][0]
+            assert "--network" in cmd
+            assert cmd[cmd.index("--network") + 1] == "none"
+
+    def test_start_container_opt_in_network(self):
+        with (
+            patch.object(cs_mod, "CONTAINER_IMAGE", "test-image:latest"),
+            patch.object(cs_mod, "CONTAINER_WORKSPACE", ""),
+            patch.object(cs_mod, "CONTAINER_NETWORK", "bridge"),
+            patch("subprocess.run", return_value=_make_proc(returncode=0)) as mock_run,
+        ):
+            cs_mod._start_container()
+            cmd = mock_run.call_args[0][0]
+            assert "--network" in cmd
+            assert cmd[cmd.index("--network") + 1] == "bridge"
+
+    def test_network_defaults_to_none_when_unset(self, monkeypatch):
+        original = os.environ.get("CONTAINER_NETWORK")
+        monkeypatch.delenv("CONTAINER_NETWORK", raising=False)
+        try:
+            reloaded = _reload_cs()
+            assert reloaded.CONTAINER_NETWORK == "none"
+        finally:
+            _restore_env_and_reload("CONTAINER_NETWORK", original)
+
+    def test_network_falls_back_to_none_when_blank(self, monkeypatch):
+        original = os.environ.get("CONTAINER_NETWORK")
+        monkeypatch.setenv("CONTAINER_NETWORK", "   ")
+        try:
+            reloaded = _reload_cs()
+            assert reloaded.CONTAINER_NETWORK == "none"
+        finally:
+            _restore_env_and_reload("CONTAINER_NETWORK", original)
+
 
 # ---------------------------------------------------------------------------
 # shell_exec tests
@@ -220,6 +289,17 @@ def test_persistent_name_differs_for_different_images(self):
                 name_b = cs_mod._persistent_name()
             assert name_a != name_b
 
+    def test_persistent_name_varies_with_network(self):
+        with (
+            patch.object(cs_mod, "CONTAINER_IMAGE", "test-image:latest"),
+            patch.object(cs_mod, "CONTAINER_PERSIST_KEY", ""),
+        ):
+            with patch.object(cs_mod, "CONTAINER_NETWORK", "none"):
+                name_none = cs_mod._persistent_name()
+            with patch.object(cs_mod, "CONTAINER_NETWORK", "bridge"):
+                name_bridge = cs_mod._persistent_name()
+            assert name_none != name_bridge
+
     def test_start_reuses_running_persistent_container(self):
         inspect_proc = _make_proc(
             returncode=0,
@@ -290,6 +370,63 @@ def test_is_running_returns_false_on_bad_json(self):
             assert cs_mod._is_running("test-name") is False
 
 
+# ---------------------------------------------------------------------------
+# Transport selection tests
+# ---------------------------------------------------------------------------
+
+class TestRunServer:
+    def test_default_transport_is_stdio(self, monkeypatch):
+        original = os.environ.get("CONTAINER_SHELL_TRANSPORT")
+        monkeypatch.delenv("CONTAINER_SHELL_TRANSPORT", raising=False)
+        try:
+            reloaded = _reload_cs()
+            assert reloaded.CONTAINER_SHELL_TRANSPORT == "stdio"
+        finally:
+            _restore_env_and_reload("CONTAINER_SHELL_TRANSPORT", original)
+
+    def test_blank_transport_falls_back_to_stdio(self, monkeypatch):
+        original = os.environ.get("CONTAINER_SHELL_TRANSPORT")
+        monkeypatch.setenv("CONTAINER_SHELL_TRANSPORT", "   ")
+        try:
+            reloaded = _reload_cs()
+            assert reloaded.CONTAINER_SHELL_TRANSPORT == "stdio"
+        finally:
+            _restore_env_and_reload("CONTAINER_SHELL_TRANSPORT", original)
+
+    def test_run_server_stdio_uses_stdio_call(self):
+        with (
+            patch.object(cs_mod, "CONTAINER_SHELL_TRANSPORT", "stdio"),
+            patch.object(cs_mod.mcp, "run") as mock_run,
+        ):
+            cs_mod._run_server()
+            mock_run.assert_called_once_with(show_banner=False)
+
+    @pytest.mark.parametrize("transport", ["http", "streamable-http", "sse"])
+    def test_run_server_network_transport_binds_host_port(self, transport):
+        with (
+            patch.object(cs_mod, "CONTAINER_SHELL_TRANSPORT", transport),
+            patch.object(cs_mod, "CONTAINER_SHELL_HOST", "127.0.0.1"),
+            patch.object(cs_mod, "CONTAINER_SHELL_PORT", 9123),
+            patch.object(cs_mod.mcp, "run") as mock_run,
+        ):
+            cs_mod._run_server()
+            mock_run.assert_called_once_with(
+                transport=transport,
+                host="127.0.0.1",
+                port=9123,
+                show_banner=False,
+            )
+
+    def test_run_server_rejects_unknown_transport(self):
+        with (
+            patch.object(cs_mod, "CONTAINER_SHELL_TRANSPORT", "carrier-pigeon"),
+            patch.object(cs_mod.mcp, "run") as mock_run,
+        ):
+            with pytest.raises(ValueError, match="Unsupported CONTAINER_SHELL_TRANSPORT"):
+                cs_mod._run_server()
+            mock_run.assert_not_called()
+
+
 # ---------------------------------------------------------------------------
 # Toolbox YAML validation
 # ---------------------------------------------------------------------------