sand-boxer/src/sandboxer/extensions/vm_packer.py

"""ext.vm-packer — attach to pre-built VMs or trigger Packer builds.

**attach** mode: connect via SSH to an existing VM (tunnel alias, localhost:port,
or direct host). Creates an isolated workspace directory; teardown removes the
workspace, not the VM.

**build** mode: run `packer init` + `packer build` for a template directory
(the-custodian build-machines conventions). Teardown preserves the OVA artifact.
"""

from __future__ import annotations

import os
import subprocess
from pathlib import Path
from typing import Any

from sandboxer.extensions.base import SandboxExtension
from sandboxer.extensions.ssh import SSHConfig
from sandboxer.lifecycle.state_hub import emit_progress_note
from sandboxer.models import Profile

BUILD_PROFILE_IDS = frozenset({"profile.vm-packer-build"})


class VMPackerExtension(SandboxExtension):
    """Attach sandbox workspace on a pre-provisioned VM, or run Packer build."""

    def __init__(self, config: dict[str, Any] | None = None) -> None:
        super().__init__(config)
        self.workspace_base: str = self.config.get("workspace_base", "/build")
        self.default_user: str | None = self.config.get("ssh_user")
        self.ready_timeout_s: int = int(self.config.get("ready_timeout_s", 30))
        self.default_packer_template: str | None = self.config.get("default_packer_template")
        self.packer_bin: str = self.config.get("packer_bin", "packer")

    def _resolve_mode(self, profile: Profile, inputs: dict[str, str]) -> str:
        mode = inputs.get("mode")
        if mode:
            return mode
        if profile.id in BUILD_PROFILE_IDS:
            return "build"
        return "attach"

    def _ssh_from_handle(self, handle: dict[str, str]) -> SSHConfig:
        port_raw = handle.get("ssh_port") or os.environ.get("SANDBOXER_VM_SSH_PORT")
        port = int(port_raw) if port_raw else None
        user = handle.get("ssh_user") or self.default_user
        host = handle.get("vm_host") or handle.get("host") or "localhost"
        return SSHConfig(
            host=host,
            user=user,
            key=os.environ.get("SANDBOXER_SSH_KEY"),
            port=port,
        )

    def _resolve_vm_host(
        self, inputs: dict[str, str], placement_host: str
    ) -> tuple[str, int | None]:
        """Return SSH host and optional port for attach mode."""
        if inputs.get("ssh_port"):
            return inputs.get("vm_host") or placement_host or "localhost", int(inputs["ssh_port"])
        if inputs.get("tunnel_port"):
            return "localhost", int(inputs["tunnel_port"])
        env_port = os.environ.get("SANDBOXER_VM_TUNNEL_PORT")
        if env_port and (placement_host in ("localhost", "127.0.0.1", "workstation")):
            return "localhost", int(env_port)
        return inputs.get("vm_host") or placement_host, None

    def _run_packer(
        self,
        args: list[str],
        *,
        cwd: Path,
        sandbox_id: str,
        profile_id: str,
    ) -> None:
        cmd = [self.packer_bin, *args]
        emit_progress_note(
            f"Packer: {' '.join(cmd)}",
            sandbox_id=sandbox_id,
            profile_id=profile_id,
            detail={"cwd": str(cwd)},
        )
        result = subprocess.run(
            cmd,
            cwd=cwd,
            capture_output=True,
            text=True,
            check=False,
        )
        if result.returncode != 0:
            msg = result.stderr.strip() or result.stdout.strip() or "unknown error"
            raise RuntimeError(f"Packer command failed ({' '.join(cmd)}): {msg}")

    def _find_ova_artifact(self, template_dir: Path, vm_name: str) -> Path | None:
        candidates = sorted(
            template_dir.glob(f"{vm_name}*.ova"),
            key=lambda p: p.stat().st_mtime,
            reverse=True,
        )
        if candidates:
            return candidates[0]
        all_ova = sorted(
            template_dir.glob("*.ova"),
            key=lambda p: p.stat().st_mtime,
            reverse=True,
        )
        return all_ova[0] if all_ova else None

    def _provision_build(
        self, profile: Profile, inputs: dict[str, str], host: str
    ) -> dict[str, str]:
        sandbox_id = self.new_sandbox_id(inputs)
        template_raw = (
            inputs.get("packer_template")
            or os.environ.get("SANDBOXER_PACKER_TEMPLATE")
            or self.default_packer_template
        )
        if not template_raw:
            raise ValueError(
                "inputs.packer_template is required for build mode "
                "(or set SANDBOXER_PACKER_TEMPLATE)"
            )

        vm_name = inputs.get("vm_name") or inputs.get("vm")
        if not vm_name:
            raise ValueError("inputs.vm_name or inputs.vm is required for build mode")

        template_dir = Path(template_raw).expanduser().resolve()
        if not template_dir.is_dir():
            raise FileNotFoundError(f"Packer template directory not found: {template_dir}")

        pkr_files = list(template_dir.glob("*.pkr.hcl")) + list(template_dir.glob("*.pkr.json"))
        if not pkr_files:
            raise FileNotFoundError(f"No Packer template (*.pkr.hcl) in {template_dir}")

        self._run_packer(
            ["init", "."],
            cwd=template_dir,
            sandbox_id=sandbox_id,
            profile_id=profile.id,
        )
        build_args = ["build", "-var", f"vm_name={vm_name}"]
        for key, value in inputs.items():
            if key.startswith("packer_var_"):
                build_args.extend(["-var", f"{key.removeprefix('packer_var_')}={value}"])
        self._run_packer(
            build_args,
            cwd=template_dir,
            sandbox_id=sandbox_id,
            profile_id=profile.id,
        )

        artifact = self._find_ova_artifact(template_dir, vm_name)
        artifact_path = str(artifact) if artifact else ""

        return {
            "sandbox_id": sandbox_id,
            "host": host,
            "mode": "build",
            "vm_name": vm_name,
            "vm_target": vm_name,
            "packer_template": str(template_dir),
            "artifact_path": artifact_path,
            "remote_dir": artifact_path,
        }

    def _provision_attach(
        self, profile: Profile, inputs: dict[str, str], host: str
    ) -> dict[str, str]:
        vm_target = inputs.get("vm") or inputs.get("ssh_target")
        if not vm_target:
            raise ValueError("inputs.vm or inputs.ssh_target is required for ext.vm-packer")

        sandbox_id = self.new_sandbox_id(inputs)
        remote_dir = inputs.get("workspace_dir") or f"{self.workspace_base}/sbx-{sandbox_id}"

        vm_host, ssh_port = self._resolve_vm_host(inputs, host)
        ssh_user = inputs.get("ssh_user") or self.default_user

        if "@" in vm_target or "." not in vm_target:
            ssh = SSHConfig(host=vm_target, user=ssh_user, port=ssh_port)
            connect_host = vm_target
        else:
            ssh = SSHConfig(host=vm_host, user=ssh_user, port=ssh_port)
            connect_host = vm_host

        rc, out = ssh.run(f"mkdir -p {remote_dir}")
        if rc != 0:
            raise RuntimeError(f"Failed to create workspace on VM: {out}")

        repo_path_str = ""
        repo = inputs.get("repo")
        if repo:
            repo_path = Path(repo).expanduser().resolve()
            if not repo_path.exists():
                raise FileNotFoundError(f"Repo path does not exist: {repo_path}")
            ssh.rsync(repo_path, remote_dir)
            repo_path_str = str(repo_path)

        return {
            "sandbox_id": sandbox_id,
            "host": host,
            "vm_host": connect_host,
            "vm_target": vm_target,
            "remote_dir": remote_dir,
            "ssh_user": ssh.user or "",
            "ssh_port": str(ssh.port) if ssh.port else "",
            "mode": "attach",
            "repo": repo_path_str,
        }

    def provision(
        self, profile: Profile, inputs: dict[str, str], host: str
    ) -> dict[str, str]:
        mode = self._resolve_mode(profile, inputs)
        if mode == "build":
            return self._provision_build(profile, inputs, host)
        return self._provision_attach(profile, inputs, host)

    def wait_ready(self, handle: dict[str, str]) -> dict[str, str]:
        if handle.get("mode") == "build":
            artifact = handle.get("artifact_path") or handle.get("remote_dir")
            if not artifact or not Path(artifact).is_file():
                raise RuntimeError(f"Packer artifact not found: {artifact}")
            return {
                "host": handle.get("host"),
                "remote_dir": artifact,
                "endpoint": artifact,
            }

        ssh = self._ssh_from_handle(handle)
        remote_dir = handle["remote_dir"]
        cmd = f"test -d {remote_dir} && echo ready"
        rc, out = ssh.run(cmd, timeout=self.ready_timeout_s)
        if rc != 0 or "ready" not in out:
            raise RuntimeError(f"VM workspace not ready: {out}")
        return {
            "ssh": ssh.target,
            "remote_dir": remote_dir,
            "host": handle.get("host"),
        }

    def teardown(self, handle: dict[str, str]) -> dict[str, str]:
        if handle.get("mode") == "build":
            artifact = handle.get("artifact_path") or handle.get("remote_dir") or ""
            return {
                "workspace_removed": "false",
                "remote_dir": artifact,
                "vm_preserved": "true",
                "artifact_preserved": "true",
            }

        remote_dir = handle.get("remote_dir")
        cleaned_dir = False
        if remote_dir:
            ssh = self._ssh_from_handle(handle)
            rc, _ = ssh.run(f"rm -rf {remote_dir}", timeout=60)
            cleaned_dir = rc == 0
        return {
            "workspace_removed": str(cleaned_dir),
            "remote_dir": remote_dir or "",
            "vm_preserved": "true",
        }