from __future__ import annotations import subprocess import uuid from pathlib import Path from .config import UPLOAD_DIR from .log import get_logger _LOG = get_logger(__name__) class LinuxVM: """Manage a lightweight Linux VM using Docker.""" def __init__( self, image: str = "ubuntu:latest", host_dir: str = UPLOAD_DIR ) -> None: self._image = image self._name = f"chat-vm-{uuid.uuid4().hex[:8]}" self._running = False self._host_dir = Path(host_dir) self._host_dir.mkdir(parents=True, exist_ok=True) def start(self) -> None: """Start the VM if it is not already running.""" if self._running: return try: subprocess.run( ["docker", "pull", self._image], check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) subprocess.run( [ "docker", "run", "-d", "--name", self._name, "-v", f"{self._host_dir}:/data", self._image, "sleep", "infinity", ], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) self._running = True except Exception as exc: # pragma: no cover - runtime failures _LOG.error("Failed to start VM: %s", exc) raise RuntimeError(f"Failed to start VM: {exc}") from exc def execute(self, command: str, *, timeout: int = 3) -> str: """Execute a command inside the running VM.""" if not self._running: raise RuntimeError("VM is not running") try: completed = subprocess.run( [ "docker", "exec", self._name, "bash", "-lc", command, ], capture_output=True, text=True, timeout=timeout, ) except subprocess.TimeoutExpired as exc: return f"Command timed out after {timeout}s: {exc.cmd}" except Exception as exc: # pragma: no cover - unforeseen errors return f"Failed to execute command: {exc}" output = completed.stdout if completed.stderr: output = f"{output}\n{completed.stderr}" if output else completed.stderr return output.strip() def stop(self) -> None: """Terminate the VM if running.""" if not self._running: return subprocess.run( ["docker", "rm", "-f", self._name], check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) self._running = False def __enter__(self) -> "LinuxVM": self.start() return self def __exit__(self, exc_type, exc, tb) -> None: self.stop()