strix-agent 0.1.9__py3-none-any.whl → 0.1.11__py3-none-any.whl

strix/runtime/docker_runtime.py

@@ -7,19 +7,15 @@ from pathlib import Path
 from typing import cast
 
 import docker
-from docker.errors import DockerException, NotFound
+from docker.errors import DockerException, ImageNotFound, NotFound
 from docker.models.containers import Container
 
 from .runtime import AbstractRuntime, SandboxInfo
 
 
-STRIX_AGENT_LABEL = "StrixAgent_ID"
-STRIX_SCAN_LABEL = "StrixScan_ID"
-STRIX_IMAGE = os.getenv("STRIX_IMAGE", "ghcr.io/usestrix/strix-sandbox:0.1.4")
+STRIX_IMAGE = os.getenv("STRIX_IMAGE", "ghcr.io/usestrix/strix-sandbox:0.1.10")
 logger = logging.getLogger(__name__)
 
-_initialized_volumes: set[str] = set()
-
 
 class DockerRuntime(AbstractRuntime):
     def __init__(self) -> None:
@@ -29,9 +25,18 @@ class DockerRuntime(AbstractRuntime):
             logger.exception("Failed to connect to Docker daemon")
             raise RuntimeError("Docker is not available or not configured correctly.") from e
 
+        self._scan_container: Container | None = None
+        self._tool_server_port: int | None = None
+        self._tool_server_token: str | None = None
+
     def _generate_sandbox_token(self) -> str:
         return secrets.token_urlsafe(32)
 
+    def _find_available_port(self) -> int:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            return cast("int", s.getsockname()[1])
+
     def _get_scan_id(self, agent_id: str) -> str:
         try:
             from strix.cli.tracer import get_global_tracer
@@ -46,37 +51,151 @@ class DockerRuntime(AbstractRuntime):
 
         return f"scan-{agent_id.split('-')[0]}"
 
-    def _find_available_port(self) -> int:
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-            s.bind(("", 0))
-            return cast("int", s.getsockname()[1])
+    def _verify_image_available(self, image_name: str, max_retries: int = 3) -> None:
+        def _validate_image(image: docker.models.images.Image) -> None:
+            if not image.id or not image.attrs:
+                raise ImageNotFound(f"Image {image_name} metadata incomplete")
+
+        for attempt in range(max_retries):
+            try:
+                image = self.client.images.get(image_name)
+                _validate_image(image)
+            except ImageNotFound:
+                if attempt == max_retries - 1:
+                    logger.exception(f"Image {image_name} not found after {max_retries} attempts")
+                    raise
+                logger.warning(f"Image {image_name} not ready, attempt {attempt + 1}/{max_retries}")
+                time.sleep(2**attempt)
+            except DockerException:
+                if attempt == max_retries - 1:
+                    logger.exception(f"Failed to verify image {image_name}")
+                    raise
+                logger.warning(f"Docker error verifying image, attempt {attempt + 1}/{max_retries}")
+                time.sleep(2**attempt)
+            else:
+                logger.debug(f"Image {image_name} verified as available")
+                return
 
-    def _get_workspace_volume_name(self, scan_id: str) -> str:
-        return f"strix-workspace-{scan_id}"
+    def _create_container_with_retry(self, scan_id: str, max_retries: int = 3) -> Container:
+        last_exception = None
 
-    def _get_sandbox_by_agent_id(self, agent_id: str) -> Container | None:
-        try:
-            containers = self.client.containers.list(
-                filters={"label": f"{STRIX_AGENT_LABEL}={agent_id}"}
-            )
-            if not containers:
-                return None
-            if len(containers) > 1:
-                logger.warning(
-                    "Multiple sandboxes found for agent ID %s, using the first one.", agent_id
+        for attempt in range(max_retries):
+            try:
+                self._verify_image_available(STRIX_IMAGE)
+
+                caido_port = self._find_available_port()
+                tool_server_port = self._find_available_port()
+                tool_server_token = self._generate_sandbox_token()
+
+                self._tool_server_port = tool_server_port
+                self._tool_server_token = tool_server_token
+
+                container = self.client.containers.run(
+                    STRIX_IMAGE,
+                    command="sleep infinity",
+                    detach=True,
+                    name=f"strix-scan-{scan_id}",
+                    hostname=f"strix-scan-{scan_id}",
+                    ports={
+                        f"{caido_port}/tcp": caido_port,
+                        f"{tool_server_port}/tcp": tool_server_port,
+                    },
+                    cap_add=["NET_ADMIN", "NET_RAW"],
+                    labels={"strix-scan-id": scan_id},
+                    environment={
+                        "PYTHONUNBUFFERED": "1",
+                        "CAIDO_PORT": str(caido_port),
+                        "TOOL_SERVER_PORT": str(tool_server_port),
+                        "TOOL_SERVER_TOKEN": tool_server_token,
+                    },
+                    tty=True,
                 )
-            return cast("Container", containers[0])
-        except DockerException as e:
-            logger.warning("Failed to get sandbox by agent ID %s: %s", agent_id, e)
-            return None
 
-    def _ensure_workspace_volume(self, volume_name: str) -> None:
+                self._scan_container = container
+                logger.info("Created container %s for scan %s", container.id, scan_id)
+
+                self._initialize_container(
+                    container, caido_port, tool_server_port, tool_server_token
+                )
+            except DockerException as e:
+                last_exception = e
+                if attempt == max_retries - 1:
+                    logger.exception(f"Failed to create container after {max_retries} attempts")
+                    break
+
+                logger.warning(f"Container creation attempt {attempt + 1}/{max_retries} failed")
+
+                self._tool_server_port = None
+                self._tool_server_token = None
+
+                sleep_time = (2**attempt) + (0.1 * attempt)
+                time.sleep(sleep_time)
+            else:
+                return container
+
+        raise RuntimeError(
+            f"Failed to create Docker container after {max_retries} attempts: {last_exception}"
+        ) from last_exception
+
+    def _get_or_create_scan_container(self, scan_id: str) -> Container:
+        if self._scan_container:
+            try:
+                self._scan_container.reload()
+                if self._scan_container.status == "running":
+                    return self._scan_container
+            except NotFound:
+                self._scan_container = None
+                self._tool_server_port = None
+                self._tool_server_token = None
+
         try:
-            self.client.volumes.get(volume_name)
-            logger.info(f"Using existing workspace volume: {volume_name}")
-        except NotFound:
-            self.client.volumes.create(name=volume_name, driver="local")
-            logger.info(f"Created new workspace volume: {volume_name}")
+            containers = self.client.containers.list(filters={"label": f"strix-scan-id={scan_id}"})
+            if containers:
+                container = cast("Container", containers[0])
+                if container.status != "running":
+                    container.start()
+                    time.sleep(2)
+                self._scan_container = container
+
+                for env_var in container.attrs["Config"]["Env"]:
+                    if env_var.startswith("TOOL_SERVER_PORT="):
+                        self._tool_server_port = int(env_var.split("=")[1])
+                    elif env_var.startswith("TOOL_SERVER_TOKEN="):
+                        self._tool_server_token = env_var.split("=")[1]
+
+                return container
+        except DockerException as e:
+            logger.warning("Failed to find existing container for scan %s: %s", scan_id, e)
+
+        logger.info("Creating new Docker container for scan %s", scan_id)
+        return self._create_container_with_retry(scan_id)
+
+    def _initialize_container(
+        self, container: Container, caido_port: int, tool_server_port: int, tool_server_token: str
+    ) -> None:
+        logger.info("Initializing Caido proxy on port %s", caido_port)
+        result = container.exec_run(
+            f"bash -c 'export CAIDO_PORT={caido_port} && /usr/local/bin/docker-entrypoint.sh true'",
+            detach=False,
+        )
+
+        time.sleep(5)
+
+        result = container.exec_run(
+            "bash -c 'source /etc/profile.d/proxy.sh && echo $CAIDO_API_TOKEN'", user="pentester"
+        )
+        caido_token = result.output.decode().strip() if result.exit_code == 0 else ""
+
+        container.exec_run(
+            f"bash -c 'source /etc/profile.d/proxy.sh && cd /app && "
+            f"STRIX_SANDBOX_MODE=true CAIDO_API_TOKEN={caido_token} CAIDO_PORT={caido_port} "
+            f"poetry run python strix/runtime/tool_server.py --token {tool_server_token} "
+            f"--host 0.0.0.0 --port {tool_server_port} &'",
+            detach=True,
+            user="pentester",
+        )
+
+        time.sleep(5)
 
     def _copy_local_directory_to_container(self, container: Container, local_path: str) -> None:
         import tarfile
@@ -85,10 +204,10 @@ class DockerRuntime(AbstractRuntime):
         try:
             local_path_obj = Path(local_path).resolve()
             if not local_path_obj.exists() or not local_path_obj.is_dir():
-                logger.warning(f"Local path does not exist or is not a directory: {local_path_obj}")
+                logger.warning(f"Local path does not exist or is not directory: {local_path_obj}")
                 return
 
-            logger.info(f"Copying local directory {local_path_obj} to container {container.id}")
+            logger.info(f"Copying local directory {local_path_obj} to container")
 
             tar_buffer = BytesIO()
             with tarfile.open(fileobj=tar_buffer, mode="w") as tar:
@@ -98,18 +217,14 @@ class DockerRuntime(AbstractRuntime):
                         tar.add(item, arcname=arcname)
 
             tar_buffer.seek(0)
-
-            container.put_archive("/shared_workspace", tar_buffer.getvalue())
+            container.put_archive("/workspace", tar_buffer.getvalue())
 
             container.exec_run(
-                "chown -R pentester:pentester /shared_workspace && chmod -R 755 /shared_workspace",
+                "chown -R pentester:pentester /workspace && chmod -R 755 /workspace",
                 user="root",
             )
 
-            logger.info(
-                f"Successfully copied {local_path_obj} to /shared_workspace in container "
-                f"{container.id}"
-            )
+            logger.info("Successfully copied local directory to /workspace")
 
         except (OSError, DockerException):
             logger.exception("Failed to copy local directory to container")
@@ -117,94 +232,56 @@ class DockerRuntime(AbstractRuntime):
     async def create_sandbox(
         self, agent_id: str, existing_token: str | None = None, local_source_path: str | None = None
     ) -> SandboxInfo:
-        sandbox = self._get_sandbox_by_agent_id(agent_id)
-        auth_token = existing_token or self._generate_sandbox_token()
-
         scan_id = self._get_scan_id(agent_id)
-        volume_name = self._get_workspace_volume_name(scan_id)
-
-        self._ensure_workspace_volume(volume_name)
-
-        if not sandbox:
-            logger.info("Creating new Docker sandbox for agent %s", agent_id)
-            try:
-                tool_server_port = self._find_available_port()
-                caido_port = self._find_available_port()
+        container = self._get_or_create_scan_container(scan_id)
 
-                volumes_config = {volume_name: {"bind": "/shared_workspace", "mode": "rw"}}
-                container_name = f"strix-{agent_id}"
-
-                sandbox = self.client.containers.run(
-                    STRIX_IMAGE,
-                    command="sleep infinity",
-                    detach=True,
-                    name=container_name,
-                    hostname=container_name,
-                    ports={
-                        f"{tool_server_port}/tcp": tool_server_port,
-                        f"{caido_port}/tcp": caido_port,
-                    },
-                    cap_add=["NET_ADMIN", "NET_RAW"],
-                    labels={
-                        STRIX_AGENT_LABEL: agent_id,
-                        STRIX_SCAN_LABEL: scan_id,
-                    },
-                    environment={
-                        "PYTHONUNBUFFERED": "1",
-                        "STRIX_AGENT_ID": agent_id,
-                        "STRIX_SANDBOX_TOKEN": auth_token,
-                        "STRIX_TOOL_SERVER_PORT": str(tool_server_port),
-                        "CAIDO_PORT": str(caido_port),
-                    },
-                    volumes=volumes_config,
-                    tty=True,
-                )
-                logger.info(
-                    "Created new sandbox %s for agent %s with shared workspace %s",
-                    sandbox.id,
-                    agent_id,
-                    volume_name,
-                )
-            except DockerException as e:
-                raise RuntimeError(f"Failed to create Docker sandbox: {e}") from e
+        source_copied_key = f"_source_copied_{scan_id}"
+        if local_source_path and not hasattr(self, source_copied_key):
+            self._copy_local_directory_to_container(container, local_source_path)
+            setattr(self, source_copied_key, True)
 
-        assert sandbox is not None
-        if sandbox.status != "running":
-            sandbox.start()
-            time.sleep(15)
+        container_id = container.id
+        if container_id is None:
+            raise RuntimeError("Docker container ID is unexpectedly None")
 
-        if local_source_path and volume_name not in _initialized_volumes:
-            self._copy_local_directory_to_container(sandbox, local_source_path)
-            _initialized_volumes.add(volume_name)
+        token = existing_token if existing_token is not None else self._tool_server_token
 
-        sandbox_id = sandbox.id
-        if sandbox_id is None:
-            raise RuntimeError("Docker container ID is unexpectedly None")
+        if self._tool_server_port is None or token is None:
+            raise RuntimeError("Tool server not initialized or no token available")
 
-        tool_server_port_str = sandbox.attrs["Config"]["Env"][
-            next(
-                (
-                    i
-                    for i, s in enumerate(sandbox.attrs["Config"]["Env"])
-                    if s.startswith("STRIX_TOOL_SERVER_PORT=")
-                ),
-                -1,
-            )
-        ].split("=")[1]
-        tool_server_port = int(tool_server_port_str)
+        api_url = await self.get_sandbox_url(container_id, self._tool_server_port)
 
-        api_url = await self.get_sandbox_url(sandbox_id, tool_server_port)
+        await self._register_agent_with_tool_server(api_url, agent_id, token)
 
         return {
-            "workspace_id": sandbox_id,
+            "workspace_id": container_id,
             "api_url": api_url,
-            "auth_token": auth_token,
-            "tool_server_port": tool_server_port,
+            "auth_token": token,
+            "tool_server_port": self._tool_server_port,
+            "agent_id": agent_id,
         }
 
-    async def get_sandbox_url(self, sandbox_id: str, port: int) -> str:
+    async def _register_agent_with_tool_server(
+        self, api_url: str, agent_id: str, token: str
+    ) -> None:
+        import httpx
+
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    f"{api_url}/register_agent",
+                    params={"agent_id": agent_id},
+                    headers={"Authorization": f"Bearer {token}"},
+                    timeout=30,
+                )
+                response.raise_for_status()
+                logger.info(f"Registered agent {agent_id} with tool server")
+        except (httpx.RequestError, httpx.HTTPStatusError) as e:
+            logger.warning(f"Failed to register agent {agent_id}: {e}")
+
+    async def get_sandbox_url(self, container_id: str, port: int) -> str:
         try:
-            container = self.client.containers.get(sandbox_id)
+            container = self.client.containers.get(container_id)
             container.reload()
 
             host = "localhost"
@@ -214,58 +291,25 @@ class DockerRuntime(AbstractRuntime):
                     host = docker_host.split("://")[1].split(":")[0]
 
         except NotFound:
-            raise ValueError(f"Sandbox {sandbox_id} not found.") from None
+            raise ValueError(f"Container {container_id} not found.") from None
         except DockerException as e:
-            raise RuntimeError(f"Failed to get sandbox URL for {sandbox_id}: {e}") from e
+            raise RuntimeError(f"Failed to get container URL for {container_id}: {e}") from e
         else:
             return f"http://{host}:{port}"
 
-    async def destroy_sandbox(self, sandbox_id: str) -> None:
-        logger.info("Destroying Docker sandbox %s", sandbox_id)
+    async def destroy_sandbox(self, container_id: str) -> None:
+        logger.info("Destroying scan container %s", container_id)
         try:
-            container = self.client.containers.get(sandbox_id)
-
-            scan_id = None
-            if container.labels and STRIX_SCAN_LABEL in container.labels:
-                scan_id = container.labels[STRIX_SCAN_LABEL]
-
+            container = self.client.containers.get(container_id)
             container.stop()
             container.remove()
-            logger.info("Successfully destroyed sandbox %s", sandbox_id)
+            logger.info("Successfully destroyed container %s", container_id)
 
-            if scan_id:
-                await self._cleanup_workspace_if_empty(scan_id)
+            self._scan_container = None
+            self._tool_server_port = None
+            self._tool_server_token = None
 
         except NotFound:
-            logger.warning("Sandbox %s not found for destruction.", sandbox_id)
+            logger.warning("Container %s not found for destruction.", container_id)
         except DockerException as e:
-            logger.warning("Failed to destroy sandbox %s: %s", sandbox_id, e)
-
-    async def _cleanup_workspace_if_empty(self, scan_id: str) -> None:
-        try:
-            volume_name = self._get_workspace_volume_name(scan_id)
-
-            containers = self.client.containers.list(
-                all=True, filters={"label": f"{STRIX_SCAN_LABEL}={scan_id}"}
-            )
-
-            if not containers:
-                try:
-                    volume = self.client.volumes.get(volume_name)
-                    volume.remove()
-                    logger.info(
-                        f"Cleaned up workspace volume {volume_name} for completed scan {scan_id}"
-                    )
-
-                    _initialized_volumes.discard(volume_name)
-
-                except NotFound:
-                    logger.debug(f"Volume {volume_name} already removed")
-                except DockerException as e:
-                    logger.warning(f"Failed to remove volume {volume_name}: {e}")
-
-        except DockerException as e:
-            logger.warning("Error during workspace cleanup for scan %s: %s", scan_id, e)
-
-    async def cleanup_scan_workspace(self, scan_id: str) -> None:
-        await self._cleanup_workspace_if_empty(scan_id)
+            logger.warning("Failed to destroy container %s: %s", container_id, e)

strix/runtime/runtime.py

@@ -7,6 +7,7 @@ class SandboxInfo(TypedDict):
     api_url: str
     auth_token: str | None
     tool_server_port: int
+    agent_id: str
 
 
 class AbstractRuntime(ABC):
@@ -17,9 +18,9 @@ class AbstractRuntime(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def get_sandbox_url(self, sandbox_id: str, port: int) -> str:
+    async def get_sandbox_url(self, container_id: str, port: int) -> str:
         raise NotImplementedError
 
     @abstractmethod
-    async def destroy_sandbox(self, sandbox_id: str) -> None:
+    async def destroy_sandbox(self, container_id: str) -> None:
         raise NotImplementedError