276 lines
9.6 KiB
Python
276 lines
9.6 KiB
Python
"""sandbox.py – Docker container lifecycle and execution environment."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import select as _select
|
||
import socket as _socket
|
||
import tarfile
|
||
from contextlib import ExitStack
|
||
from select import select as _original_select
|
||
from pathlib import Path
|
||
from typing import TYPE_CHECKING
|
||
from unittest.mock import patch
|
||
|
||
_original_poll = getattr(_select, "poll", None)
|
||
|
||
import docker
|
||
import docker.errors
|
||
from docker.utils.socket import consume_socket_output, demux_adaptor, frames_iter
|
||
from loguru import logger
|
||
|
||
if TYPE_CHECKING:
|
||
import docker.models.containers
|
||
|
||
|
||
class DockerSandbox:
|
||
"""
|
||
Manages a single long-running Docker container.
|
||
|
||
All container configuration (volumes, ports, environment variables,
|
||
capabilities) is supplied by the caller.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
container_name: str,
|
||
image: str = "docker-agent-sandbox",
|
||
dockerfile_dir: str | None = None,
|
||
volumes: dict | None = None,
|
||
ports: dict | None = None,
|
||
environment: dict[str, str] | None = None,
|
||
working_dir: str | None = None,
|
||
# TODO(security): default network_mode is "bridge", giving containers unrestricted
|
||
# outbound network access. For an untrusted-code sandbox this should default to
|
||
# "none" so callers must explicitly opt in to network access.
|
||
network_mode: str = "bridge",
|
||
cap_drop: list[str] | None = None,
|
||
cap_add: list[str] | None = None,
|
||
security_opt: list[str] | None = None,
|
||
cpu_limit: float = 8,
|
||
memory_limit: str = "16g",
|
||
command: str | None = None,
|
||
) -> None:
|
||
self.container_name = container_name
|
||
self._image = image
|
||
self._dockerfile_dir = dockerfile_dir
|
||
self._volumes = volumes or {}
|
||
self._ports = ports or {}
|
||
self._environment = environment or {}
|
||
self._working_dir = working_dir
|
||
self._network_mode = network_mode
|
||
self._cap_drop = cap_drop
|
||
self._cap_add = cap_add
|
||
self._security_opt = security_opt
|
||
self._nano_cpus = int(cpu_limit * 1e9)
|
||
self._memory_limit = memory_limit
|
||
self._command = command
|
||
self._client: docker.DockerClient = docker.from_env()
|
||
self._container: docker.models.containers.Container | None = None
|
||
|
||
def __enter__(self) -> "DockerSandbox":
|
||
return self
|
||
|
||
def __exit__(self, *_: object) -> None:
|
||
self.stop()
|
||
self._client.close()
|
||
|
||
# ------------------------------------------------------------------
|
||
# Lifecycle
|
||
# ------------------------------------------------------------------
|
||
|
||
def build_image_if_missing(self) -> None:
|
||
"""Build the Docker image if it is not already present locally."""
|
||
try:
|
||
self._client.images.get(self._image)
|
||
logger.info("Image {!r} already present, skipping build.", self._image)
|
||
return
|
||
except docker.errors.ImageNotFound:
|
||
pass
|
||
|
||
if self._dockerfile_dir is None:
|
||
raise ValueError("dockerfile_dir must be provided to build the image")
|
||
|
||
logger.info("Building image {!r} from {} …", self._image, self._dockerfile_dir)
|
||
_, logs = self._client.images.build(
|
||
path=self._dockerfile_dir,
|
||
tag=self._image,
|
||
rm=True,
|
||
)
|
||
for entry in logs:
|
||
line = entry.get("stream", "").rstrip()
|
||
if line:
|
||
logger.debug(" {}", line)
|
||
|
||
logger.success("Image {!r} built successfully.", self._image)
|
||
|
||
def start(self) -> None:
|
||
"""
|
||
Start the container.
|
||
|
||
Any existing container with the same name is removed first so that
|
||
re-runs always start from a clean state.
|
||
"""
|
||
try:
|
||
old = self._client.containers.get(self.container_name)
|
||
old.remove(force=True)
|
||
except docker.errors.NotFound:
|
||
pass
|
||
|
||
run_kwargs: dict = dict(
|
||
name=self.container_name,
|
||
detach=True,
|
||
volumes=self._volumes,
|
||
environment=self._environment,
|
||
ports=self._ports,
|
||
network_mode=self._network_mode,
|
||
)
|
||
if self._working_dir is not None:
|
||
run_kwargs["working_dir"] = self._working_dir
|
||
if self._cap_drop is not None:
|
||
run_kwargs["cap_drop"] = self._cap_drop
|
||
if self._cap_add is not None:
|
||
run_kwargs["cap_add"] = self._cap_add
|
||
if self._security_opt is not None:
|
||
run_kwargs["security_opt"] = self._security_opt
|
||
run_kwargs["nano_cpus"] = self._nano_cpus
|
||
run_kwargs["mem_limit"] = self._memory_limit
|
||
|
||
try:
|
||
self._container = self._client.containers.run(self._image, self._command, **run_kwargs)
|
||
except docker.errors.ImageNotFound:
|
||
raise RuntimeError(
|
||
f"Image {self._image!r} not found locally. "
|
||
"Call build_image_if_missing() first, or pull/build the image manually."
|
||
) from None
|
||
self._container.reload()
|
||
|
||
logger.info(
|
||
"Container {!r} started (id={}).",
|
||
self.container_name,
|
||
self._container.short_id,
|
||
)
|
||
|
||
def stop(self) -> None:
|
||
"""Remove the container."""
|
||
if self._container is not None:
|
||
self._container.remove(force=True)
|
||
logger.info("Container {!r} stopped.", self.container_name)
|
||
self._container = None
|
||
|
||
# ------------------------------------------------------------------
|
||
# Port inspection
|
||
# ------------------------------------------------------------------
|
||
|
||
def get_host_port(self, container_port: str) -> str:
|
||
"""
|
||
Return the host port Docker mapped to *container_port*.
|
||
|
||
*container_port* should include the protocol, e.g. ``"8080/tcp"``.
|
||
Raises ``RuntimeError`` if the container is not running or the port
|
||
is not exposed.
|
||
"""
|
||
if self._container is None:
|
||
raise RuntimeError("Sandbox container is not running.")
|
||
mappings = self._container.ports.get(container_port)
|
||
if not mappings:
|
||
raise RuntimeError(
|
||
f"Port {container_port!r} is not mapped on container {self.container_name!r}"
|
||
)
|
||
return mappings[0]["HostPort"]
|
||
|
||
# ------------------------------------------------------------------
|
||
# File I/O
|
||
# ------------------------------------------------------------------
|
||
|
||
def _resolve_path(self, path: str) -> Path:
|
||
"""Resolve *path* against working_dir if relative."""
|
||
p = Path(path)
|
||
if not p.is_absolute() and self._working_dir is not None:
|
||
p = Path(self._working_dir) / p
|
||
return p
|
||
|
||
def read_file(self, path: str) -> bytes:
|
||
"""
|
||
Read the file at *path* from the container using ``get_archive``.
|
||
|
||
Returns the raw file bytes. Raises ``FileNotFoundError`` if the path
|
||
does not exist and ``RuntimeError`` if the container is not running.
|
||
"""
|
||
if self._container is None:
|
||
raise RuntimeError("Sandbox container is not running.")
|
||
|
||
p = self._resolve_path(path)
|
||
|
||
try:
|
||
stream, _ = self._container.get_archive(str(p))
|
||
except docker.errors.NotFound:
|
||
raise FileNotFoundError(f"No such file in container: {path!r}") from None
|
||
|
||
buf = io.BytesIO()
|
||
for chunk in stream:
|
||
buf.write(chunk)
|
||
buf.seek(0)
|
||
|
||
with tarfile.open(fileobj=buf) as tar:
|
||
member = tar.getmembers()[0]
|
||
f = tar.extractfile(member)
|
||
if f is None:
|
||
raise IsADirectoryError(f"{path!r} is a directory, not a file")
|
||
return f.read()
|
||
|
||
def write_file(self, path: str, content: str) -> None:
|
||
"""
|
||
Write *content* to *path* inside the container using ``put_archive``.
|
||
|
||
The tar archive API avoids all shell-escaping concerns. Parent
|
||
directories are created automatically via a preceding ``mkdir -p``.
|
||
"""
|
||
if self._container is None:
|
||
raise RuntimeError("Sandbox container is not running.")
|
||
|
||
p = self._resolve_path(path)
|
||
encoded = content.encode("utf-8")
|
||
|
||
self._container.exec_run(["mkdir", "-p", str(p.parent)])
|
||
|
||
buf = io.BytesIO()
|
||
with tarfile.open(fileobj=buf, mode="w") as tar:
|
||
info = tarfile.TarInfo(name=p.name)
|
||
info.size = len(encoded)
|
||
info.mode = 0o644
|
||
tar.addfile(info, io.BytesIO(encoded))
|
||
buf.seek(0)
|
||
self._container.put_archive(str(p.parent), buf)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Command execution
|
||
# ------------------------------------------------------------------
|
||
|
||
def exec(self, command: str, timeout: int = 120, kill_grace: int = 5) -> tuple[int, str]:
|
||
if self._container is None:
|
||
return 1, "Sandbox container is not running."
|
||
|
||
wrapped = [
|
||
"timeout", "--kill-after", f"{kill_grace}s", f"{timeout}s",
|
||
"bash", "-c", command,
|
||
]
|
||
kwargs = {}
|
||
if self._working_dir is not None:
|
||
kwargs["workdir"] = self._working_dir
|
||
|
||
try:
|
||
exit_code, output = self._container.exec_run(
|
||
wrapped,
|
||
stdout=True,
|
||
stderr=True,
|
||
demux=False,
|
||
**kwargs,
|
||
)
|
||
text = (output or b"").decode("utf-8", errors="replace")
|
||
if exit_code in (124, 137):
|
||
return exit_code, text + f"\n[command timed out after {timeout}s]"
|
||
return exit_code, text
|
||
except Exception as exc:
|
||
return 1, f"Error running command in container: {exc}"
|