Add containerized shell toolbox

pyproject.toml

@@ -92,6 +92,7 @@ ignore = [
   "PLR1730",   # Replace `if` statement with `min()`
   "PLR2004",   # Magic value used in comparison
   "PLW0602",   # Using global for variable but no assignment is done
+  "PLW0603",   # Using the global statement to update a variable is discouraged
   "PLW1508",   # Invalid type for environment variable default
   "PLW1510",   # `subprocess.run` without explicit `check` argument
   "RET504",    # Unnecessary assignment before `return` statement
@@ -113,3 +114,9 @@ ignore = [
   "W291",      # Trailing whitespace
   "W293",      # Blank line contains whitespace
 ]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/*" = [
+  "S101",      # Use of assert (standard in pytest)
+  "SLF001",    # Private member accessed (tests legitimately access module internals)
+]

scripts/build_container_images.sh

@@ -0,0 +1,69 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: GitHub, Inc.
+# SPDX-License-Identifier: MIT
+
+# Build seclab container shell images.
+# Must be run from the root of the seclab-taskflows repository.
+# Images must be rebuilt whenever a Dockerfile changes.
+#
+# Usage: ./scripts/build_container_images.sh [base|malware|network|sast|all]
+#   default: all
+
+set -euo pipefail
+
+__dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+__root="$(cd "${__dir}/.." && pwd)"
+CONTAINERS_DIR="${__root}/src/seclab_taskflows/containers"
+
+build_base() {
+    echo "Building seclab-shell-base..."
+    docker build -t seclab-shell-base:latest "${CONTAINERS_DIR}/base/"
+}
+
+build_malware() {
+    echo "Building seclab-shell-malware-analysis..."
+    docker build -t seclab-shell-malware-analysis:latest "${CONTAINERS_DIR}/malware_analysis/"
+}
+
+build_network() {
+    echo "Building seclab-shell-network-analysis..."
+    docker build -t seclab-shell-network-analysis:latest "${CONTAINERS_DIR}/network_analysis/"
+}
+
+build_sast() {
+    echo "Building seclab-shell-sast..."
+    docker build -t seclab-shell-sast:latest "${CONTAINERS_DIR}/sast/"
+}
+
+target="${1:-all}"
+
+case "$target" in
+    base)
+        build_base
+        ;;
+    malware)
+        build_base
+        build_malware
+        ;;
+    network)
+        build_base
+        build_network
+        ;;
+    sast)
+        build_base
+        build_sast
+        ;;
+    all)
+        build_base
+        build_malware
+        build_network
+        build_sast
+        ;;
+    *)
+        echo "Unknown target: $target" >&2
+        echo "Usage: $0 [base|malware|network|sast|all]" >&2
+        exit 1
+        ;;
+esac
+
+echo "Done."

scripts/run_container_shell_demo.sh

@@ -0,0 +1,127 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: GitHub, Inc.
+# SPDX-License-Identifier: MIT
+
+# Run container shell demo taskflows.
+# Must be run from the root of the seclab-taskflows repository.
+#
+# Usage:
+#   ./scripts/run_container_shell_demo.sh base    [workspace_dir]
+#   ./scripts/run_container_shell_demo.sh malware [workspace_dir] [target_filename]
+#   ./scripts/run_container_shell_demo.sh network [workspace_dir] [capture_filename]
+#   ./scripts/run_container_shell_demo.sh sast    [workspace_dir] [target]
+#
+# If workspace_dir is omitted a temporary directory is used.
+# Requires AI_API_TOKEN to be set in the environment.
+
+set -euo pipefail
+
+__dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+__root="$(cd "${__dir}/.." && pwd)"
+
+export PATH="${__root}/.venv/bin:${PATH}"
+
+if [ -z "${AI_API_TOKEN:-}" ]; then
+    echo "AI_API_TOKEN is not set" >&2
+    exit 1
+fi
+
+demo="${1:-}"
+if [ -z "$demo" ]; then
+    echo "Usage: $0 <base|malware|network|sast> [workspace_dir] [target]" >&2
+    exit 1
+fi
+
+workspace="${2:-$(mktemp -d)}"
+mkdir -p "$workspace"
+
+case "$demo" in
+    base)
+        target="${3:-hello}"
+        if [ ! -f "${workspace}/${target}" ]; then
+            echo "Copying /bin/ls to ${workspace}/${target} as demo target"
+            cp /bin/ls "${workspace}/${target}"
+        fi
+        CONTAINER_WORKSPACE="$workspace" \
+        LOG_DIR="${__root}/logs" \
+        python -m seclab_taskflow_agent \
+            -t seclab_taskflows.taskflows.container_shell.demo_base \
+            -g target="$target"
+        ;;
+    malware)
+        target="${3:-suspicious.elf}"
+        if [ ! -f "${workspace}/${target}" ]; then
+            echo "Copying /bin/ls to ${workspace}/${target} as demo target"
+            cp /bin/ls "${workspace}/${target}"
+        fi
+        CONTAINER_WORKSPACE="$workspace" \
+        LOG_DIR="${__root}/logs" \
+        python -m seclab_taskflow_agent \
+            -t seclab_taskflows.taskflows.container_shell.demo_malware_analysis \
+            -g target="$target"
+        ;;
+    network)
+        capture="${3:-sample.pcap}"
+        if [ ! -f "${workspace}/${capture}" ]; then
+            echo "No pcap found at ${workspace}/${capture}" >&2
+            echo "Provide a pcap file or set workspace_dir to a directory containing one." >&2
+            exit 1
+        fi
+        CONTAINER_WORKSPACE="$workspace" \
+        LOG_DIR="${__root}/logs" \
+        python -m seclab_taskflow_agent \
+            -t seclab_taskflows.taskflows.container_shell.demo_network_analysis \
+            -g capture="$capture"
+        ;;
+    sast)
+        target="${3:-.}"
+        target_path="${workspace}/${target}"
+        if [ "$target" != "." ] && [ ! -d "$target_path" ] && [ ! -f "$target_path" ]; then
+            echo "No source found at ${target_path}" >&2
+            echo "Provide a source directory or file in workspace_dir." >&2
+            exit 1
+        fi
+        if [ "$target" = "." ] && [ -z "$(ls -A "$workspace" 2>/dev/null)" ]; then
+            echo "Generating demo Python source in ${workspace}"
+            cat > "${workspace}/demo.py" <<'PYEOF'
+import os
+import subprocess
+
+
+def read_config(path):
+    with open(path) as f:
+        return f.read()
+
+
+def run_command(cmd):
+    # intentional anti-pattern for demo purposes
+    return subprocess.run(cmd, shell=True, capture_output=True, text=True)
+
+
+def process_input(user_input):
+    result = run_command(f"echo {user_input}")
+    return result.stdout
+
+
+def main():
+    config = read_config("/etc/demo.conf") if os.path.exists("/etc/demo.conf") else ""
+    output = process_input("hello world")
+    print(config, output)
+
+
+if __name__ == "__main__":
+    main()
+PYEOF
+            target="demo.py"
+        fi
+        CONTAINER_WORKSPACE="$workspace" \
+        LOG_DIR="${__root}/logs" \
+        python -m seclab_taskflow_agent \
+            -t seclab_taskflows.taskflows.container_shell.demo_sast \
+            -g target="$target"
+        ;;
+    *)
+        echo "Unknown demo: $demo. Choose base, malware, network, or sast." >&2
+        exit 1
+        ;;
+esac

src/seclab_taskflows/containers/base/Dockerfile

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: GitHub, Inc.
+# SPDX-License-Identifier: MIT
+
+FROM debian:bookworm-slim
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    bash coreutils python3 python3-pip curl wget git ca-certificates \
+    file binutils xxd \
+    && rm -rf /var/lib/apt/lists/*
+WORKDIR /workspace

src/seclab_taskflows/containers/malware_analysis/Dockerfile

@@ -0,0 +1,17 @@
+# SPDX-FileCopyrightText: GitHub, Inc.
+# SPDX-License-Identifier: MIT
+
+FROM seclab-shell-base:latest
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    binwalk yara libimage-exiftool-perl checksec \
+    && rm -rf /var/lib/apt/lists/*
+# radare2 is not in Debian bookworm apt; install prebuilt deb from GitHub releases
+RUN ARCH=$(dpkg --print-architecture) \
+    && R2_TAG=$(curl -fsSL "https://api.github.com/repos/radareorg/radare2/releases/latest" \
+        | grep -o '"tag_name": *"[^"]*"' | grep -o '"[^"]*"$' | tr -d '"') \
+    && R2_VER="${R2_TAG#v}" \
+    && curl -fsSL "https://github.com/radareorg/radare2/releases/download/${R2_TAG}/radare2_${R2_VER}_${ARCH}.deb" \
+        -o /tmp/r2.deb \
+    && apt-get install -y /tmp/r2.deb \
+    && rm /tmp/r2.deb
+RUN pip3 install --no-cache-dir --break-system-packages pwntools capstone volatility3

src/seclab_taskflows/containers/network_analysis/Dockerfile

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: GitHub, Inc.
+# SPDX-License-Identifier: MIT
+
+FROM seclab-shell-base:latest
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    nmap tcpdump tshark netcat-openbsd dnsutils curl jq httpie \
+    && rm -rf /var/lib/apt/lists/*

src/seclab_taskflows/containers/sast/Dockerfile

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: GitHub, Inc.
+# SPDX-License-Identifier: MIT
+
+FROM seclab-shell-base:latest
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    universal-ctags global cscope ripgrep fd-find graphviz tree \
+    && ln -s /usr/bin/fdfind /usr/local/bin/fd \
+    && rm -rf /var/lib/apt/lists/*
+RUN pip3 install --no-cache-dir --break-system-packages semgrep pyan3

src/seclab_taskflows/mcp_servers/container_shell.py

@@ -0,0 +1,107 @@
+# SPDX-FileCopyrightText: GitHub, Inc.
+# SPDX-License-Identifier: MIT
+
+import atexit
+import logging
+import os
+import subprocess
+import uuid
+from typing import Annotated
+
+from fastmcp import FastMCP
+from pydantic import Field
+from seclab_taskflow_agent.path_utils import log_file_name
+
+logging.basicConfig(
+    level=logging.DEBUG,
+    format="%(asctime)s - %(levelname)s - %(message)s",
+    filename=log_file_name("container_shell.log"),
+    filemode="a",
+)
+
+mcp = FastMCP("ContainerShell")
+
+_container_name: str | None = None
+
+CONTAINER_IMAGE = os.environ.get("CONTAINER_IMAGE", "")
+CONTAINER_WORKSPACE = os.environ.get("CONTAINER_WORKSPACE", "")
+CONTAINER_TIMEOUT = int(os.environ.get("CONTAINER_TIMEOUT", "30"))
+
+_DEFAULT_WORKDIR = "/workspace"
+
+
+def _start_container() -> str:
+    """Start the Docker container and return its name."""
+    if not CONTAINER_IMAGE:
+        msg = "CONTAINER_IMAGE is not set — cannot start container"
+        raise RuntimeError(msg)
+    if CONTAINER_WORKSPACE and ":" in CONTAINER_WORKSPACE:
+        msg = f"CONTAINER_WORKSPACE must not contain a colon: {CONTAINER_WORKSPACE!r}"
+        raise RuntimeError(msg)
+    name = f"seclab-shell-{uuid.uuid4().hex[:8]}"
+    cmd = ["docker", "run", "-d", "--rm", "--name", name]
+    if CONTAINER_WORKSPACE:
+        cmd += ["-v", f"{CONTAINER_WORKSPACE}:/workspace"]
+    cmd += [CONTAINER_IMAGE, "tail", "-f", "/dev/null"]
+    logging.debug(f"Starting container: {' '.join(cmd)}")
+    result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
+    if result.returncode != 0:
+        msg = f"docker run failed: {result.stderr.strip()}"
+        raise RuntimeError(msg)
+    logging.debug(f"Container started: {name}")
+    return name
+
+
+def _stop_container() -> None:
+    """Stop the running container."""
+    global _container_name
+    if _container_name is None:
+        return
+    logging.debug(f"Stopping container: {_container_name}")
+    result = subprocess.run(
+        ["docker", "stop", "--time", "5", _container_name],
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        logging.warning(
+            "docker stop failed for container %s: %s",
+            _container_name,
+            result.stderr.strip(),
+        )
+    _container_name = None
+
+
+atexit.register(_stop_container)
+
+
+@mcp.tool()
+def shell_exec(
+    command: Annotated[str, Field(description="Shell command to execute inside the container")],
+    timeout: Annotated[int, Field(description="Timeout in seconds")] = CONTAINER_TIMEOUT,
+    workdir: Annotated[str, Field(description="Working directory inside the container")] = _DEFAULT_WORKDIR,
+) -> str:
+    """Execute a shell command inside the managed Docker container."""
+    global _container_name
+    if _container_name is None:
+        try:
+            _container_name = _start_container()
+        except RuntimeError as e:
+            return f"Failed to start container: {e}"
+
+    cmd = ["docker", "exec", "-w", workdir, _container_name, "bash", "-c", command]
+    logging.debug(f"Executing: {' '.join(cmd)}")
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+    except subprocess.TimeoutExpired:
+        return f"[exit code: timeout after {timeout}s]"
+
+    output = result.stdout
+    if result.stderr:
+        output += result.stderr
+    output += f"[exit code: {result.returncode}]"
+    return output
+
+
+if __name__ == "__main__":
+    mcp.run(show_banner=False)