InstallerRobot/agent/app/core/installer.py

from __future__ import annotations

import time
from pathlib import Path, PurePosixPath

from app.core.command_runner import CommandError, CommandRunner


APT_DPKG_OPTIONS = [
    "-o",
    "Dpkg::Use-Pty=0",
    "-o",
    "Dpkg::Options::=--force-confdef",
    "-o",
    "Dpkg::Options::=--force-confold",
]

APT_NONINTERACTIVE_ENV = {
    "DEBIAN_FRONTEND": "noninteractive",
    "DEBCONF_NONINTERACTIVE_SEEN": "true",
    "APT_LISTCHANGES_FRONTEND": "none",
}


def _parse_deb_control_output(output: str) -> dict[str, str]:
    metadata: dict[str, str] = {}

    for line in output.splitlines():
        key, separator, value = line.partition(":")
        if not separator:
            continue

        normalized_key = key.strip().lower()
        if normalized_key in {"package", "version", "architecture"}:
            metadata[normalized_key] = value.strip()

    return metadata


class DebInstaller:
    def __init__(self, command_runner: CommandRunner) -> None:
        self.command_runner = command_runner

    def get_deb_metadata(self, file_path: Path) -> dict[str, str]:
        result = self.command_runner.run([
            "dpkg-deb",
            "-f",
            str(file_path),
            "Package",
            "Version",
            "Architecture",
        ])
        metadata = _parse_deb_control_output(result.stdout)
        missing_fields = [
            field
            for field in ("package", "version", "architecture")
            if not metadata.get(field)
        ]

        if missing_fields:
            raise ValueError(
                "Downloaded .deb is missing metadata fields: "
                f"{', '.join(missing_fields)}"
            )

        return metadata

    def install_deb(self, file_path: Path) -> None:
        self.command_runner.run(
            [
                "apt-get",
                *APT_DPKG_OPTIONS,
                "install",
                "--yes",
                str(file_path),
            ],
            env=APT_NONINTERACTIVE_ENV,
        )

    def remove_package(self, package_name: str, purge: bool = False) -> None:
        action = "purge" if purge else "remove"
        self.command_runner.run(
            ["apt-get", *APT_DPKG_OPTIONS, action, "--yes", package_name],
            env=APT_NONINTERACTIVE_ENV,
        )

    def cleanup_after_remove(self, purge: bool = False) -> None:
        autoremove_command = ["apt-get", *APT_DPKG_OPTIONS, "autoremove", "--yes"]
        if purge:
            autoremove_command.append("--purge")
        self.command_runner.run(autoremove_command, env=APT_NONINTERACTIVE_ENV)
        self.command_runner.run(["apt-get", "clean"], env=APT_NONINTERACTIVE_ENV)

    def get_package_version(self, package_name: str) -> str | None:
        result = self.command_runner.run(["dpkg-query", "-W", "-f=${Version}", package_name])
        version = result.stdout.strip()
        return version or None

    def check_package_installed(self, package_name: str) -> bool:
        try:
            self.get_package_version(package_name)
            return True
        except Exception:
            return False


class AptInstaller(DebInstaller):
    def update_package_index(self) -> None:
        self.command_runner.run(
            ["apt-get", "update"],
            timeout=600,
            env=APT_NONINTERACTIVE_ENV,
        )

    def install_package(self, package_name: str) -> None:
        self.command_runner.run(
            [
                "apt-get",
                *APT_DPKG_OPTIONS,
                "install",
                "--yes",
                package_name,
            ],
            timeout=1200,
            env=APT_NONINTERACTIVE_ENV,
        )

    def discover_service_units(self, package_name: str) -> list[str]:
        """Return concrete systemd service units shipped by an installed package.

        The package name is already constrained by the Agent APT allowlist.  It is
        still passed as a single argv item and no shell expansion is performed.
        Template units are omitted because they cannot be meaningfully checked
        without an instance name.
        """
        try:
            result = self.command_runner.run(
                ["dpkg-query", "-L", package_name],
                timeout=30,
                log_output=False,
            )
        except CommandError:
            return []

        service_names: set[str] = set()
        for raw_path in result.stdout.splitlines():
            path = raw_path.strip()
            if not path.endswith(".service"):
                continue
            if "/systemd/system/" not in path:
                continue

            service_name = PurePosixPath(path).name
            if "@" in service_name:
                continue
            service_names.add(service_name)

        return sorted(service_names)

    def wait_for_postgresql(
        self,
        attempts: int = 6,
        delay_seconds: float = 2.0,
    ) -> None:
        last_error: CommandError | None = None
        for attempt in range(1, attempts + 1):
            try:
                self.command_runner.run(["pg_isready", "--timeout=5"], timeout=10)
                return
            except CommandError as error:
                last_error = error
                if attempt < attempts:
                    time.sleep(delay_seconds)

        raise RuntimeError("PostgreSQL did not become ready after installation") from last_error

    def is_postgresql_ready(self) -> bool:
        try:
            self.command_runner.run(["pg_isready", "--timeout=5"], timeout=10)
            return True
        except Exception:
            return False