DECNET/decnet/swarm/bundle_builder.py

"""Tarball + bootstrap construction for agent-enrollment bundles.

Pure I/O, no FastAPI dependency — independently testable.
"""
from __future__ import annotations

import io
import os
import pathlib
import tarfile
from datetime import datetime, timezone
from typing import Optional

from decnet.swarm import pki

# ---------------------------------------------------------------------------
# Include / exclude manifest
# ---------------------------------------------------------------------------

# Explicit include list — fails closed. Stray files on the master
# (dev venvs, .env files, editor scratch) cannot leak into the bundle.
_INCLUDED_ROOT_FILES: tuple[str, ...] = ("pyproject.toml",)
_INCLUDED_DIRS: tuple[str, ...] = ("decnet",)

# Subtrees of _INCLUDED_DIRS that must NOT ship (relative to repo root).
#   * decnet/web      — FastAPI master app, unused on agents.
#   * decnet/mutator  — swarm-wide respawn scheduler, master-only.
#   * decnet/profiler — rebuilds profiles against master DB, master-only.
_EXCLUDED_DECNET_SUBTREES: frozenset[str] = frozenset({
    "decnet/web",
    "decnet/mutator",
    "decnet/profiler",
})

# Agent-side systemd units. Profiler stays master-side intentionally.
_SYSTEMD_UNITS = (
    "decnet-agent", "decnet-forwarder", "decnet-engine", "decnet-updater",
    "decnet-collector", "decnet-prober", "decnet-sniffer",
)


# ---------------------------------------------------------------------------
# Path helpers
# ---------------------------------------------------------------------------

def _repo_root() -> pathlib.Path:
    # decnet/swarm/bundle_builder.py -> parents[2] = repo root.
    return pathlib.Path(__file__).resolve().parents[2]


def _templates_dir() -> pathlib.Path:
    return pathlib.Path(__file__).resolve().parents[1] / "web" / "templates"


# ---------------------------------------------------------------------------
# Filesystem walk
# ---------------------------------------------------------------------------

def _iter_included(root: pathlib.Path) -> list[tuple[pathlib.Path, str]]:
    """Return ``(full_path, arcname)`` pairs for every file the agent needs.

    Walk is pruned in-place: ``__pycache__`` and master-only subtrees are
    skipped at directory level so we never descend into them.
    """
    found: list[tuple[pathlib.Path, str]] = []

    for rel in _INCLUDED_ROOT_FILES:
        p = root / rel
        if p.is_file():
            found.append((p, rel))

    for top in _INCLUDED_DIRS:
        start = root / top
        if not start.is_dir():
            continue
        for dirpath, dirnames, filenames in os.walk(start, topdown=True, followlinks=False):
            dir_path = pathlib.Path(dirpath)
            rel_dir = dir_path.relative_to(root).as_posix()

            dirnames[:] = [
                d for d in dirnames
                if d not in ("__pycache__", "node_modules")
                and f"{rel_dir}/{d}" not in _EXCLUDED_DECNET_SUBTREES
            ]

            for fn in filenames:
                if fn.endswith((".pyc", ".pyo")):
                    continue
                full = dir_path / fn
                if full.is_symlink():
                    continue
                found.append((full, f"{rel_dir}/{fn}"))

    found.sort(key=lambda t: t[1])
    return found


# ---------------------------------------------------------------------------
# Content renderers
# ---------------------------------------------------------------------------

def _render_decnet_ini(
    master_host: str,
    host_uuid: str,
    use_ipvlan: bool = False,
    swarmctl_port: int = 8770,
) -> bytes:
    ipvlan_line = f"ipvlan = {'true' if use_ipvlan else 'false'}\n"
    return (
        "; Generated by DECNET agent-enrollment bundle.\n"
        "[decnet]\n"
        "mode = agent\n"
        "disallow-master = true\n"
        "log-directory = /var/log/decnet\n"
        f"{ipvlan_line}"
        "\n"
        "[agent]\n"
        f"master-host = {master_host}\n"
        f"swarmctl-port = {swarmctl_port}\n"
        "swarm-syslog-port = 6514\n"
        "agent-port = 8765\n"
        "agent-dir = /etc/decnet/agent\n"
        "updater-dir = /etc/decnet/updater\n"
        f"host-uuid = {host_uuid}\n"
    ).encode()


def _add_bytes(tar: tarfile.TarFile, name: str, data: bytes, mode: int = 0o644) -> None:
    info = tarfile.TarInfo(name)
    info.size = len(data)
    info.mode = mode
    info.mtime = int(datetime.now(timezone.utc).timestamp())
    tar.addfile(info, io.BytesIO(data))


def _render_systemd_unit(name: str, agent_name: str, master_host: str) -> bytes:
    tpl_path = _templates_dir() / f"{name}.service.j2"
    tpl = tpl_path.read_text()
    return (
        tpl.replace("{{ agent_name }}", agent_name)
           .replace("{{ master_host }}", master_host)
    ).encode()


def render_bootstrap(
    agent_name: str,
    master_host: str,
    tarball_url: str,
    expires_at: datetime,
    with_updater: bool,
) -> bytes:
    tpl_path = _templates_dir() / "enroll_bootstrap.sh.j2"
    tpl = tpl_path.read_text()
    now = datetime.now(timezone.utc).replace(microsecond=0).isoformat()
    rendered = (
        tpl.replace("{{ agent_name }}", agent_name)
           .replace("{{ master_host }}", master_host)
           .replace("{{ tarball_url }}", tarball_url)
           .replace("{{ generated_at }}", now)
           .replace("{{ expires_at }}", expires_at.replace(microsecond=0).isoformat())
           .replace("{{ with_updater }}", "true" if with_updater else "false")
    )
    return rendered.encode()


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def build_tarball(
    master_host: str,
    agent_name: str,
    host_uuid: str,
    issued: pki.IssuedCert,
    services_ini: Optional[str],
    updater_issued: Optional[pki.IssuedCert] = None,
    use_ipvlan: bool = False,
) -> bytes:
    """Return a gzipped tarball ready to be handed to the enrolling agent."""
    root = _repo_root()
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w:gz") as tar:
        for path, arcname in _iter_included(root):
            tar.add(path, arcname=arcname, recursive=False)

        _add_bytes(
            tar,
            "etc/decnet/decnet.ini",
            _render_decnet_ini(master_host, host_uuid, use_ipvlan),
        )
        for unit in _SYSTEMD_UNITS:
            _add_bytes(
                tar,
                f"etc/systemd/system/{unit}.service",
                _render_systemd_unit(unit, agent_name, master_host),
            )
        _add_bytes(tar, "home/.decnet/agent/ca.crt", issued.ca_cert_pem)
        _add_bytes(tar, "home/.decnet/agent/worker.crt", issued.cert_pem)
        _add_bytes(tar, "home/.decnet/agent/worker.key", issued.key_pem, mode=0o600)

        if updater_issued is not None:
            _add_bytes(tar, "home/.decnet/updater/ca.crt", updater_issued.ca_cert_pem)
            _add_bytes(tar, "home/.decnet/updater/updater.crt", updater_issued.cert_pem)
            _add_bytes(tar, "home/.decnet/updater/updater.key", updater_issued.key_pem, mode=0o600)

        if services_ini:
            _add_bytes(tar, "services.ini", services_ini.encode())

    return buf.getvalue()