setup.py

#!/usr/bin/env python3
"""
Setup and prerequisites check for ct-controller.

Validates system requirements, installs dependencies, and configures the host.

Usage:
    ./scripts/setup.py              # Check prerequisites (no changes)
    ./scripts/setup.py --install    # Install/configure as needed
    ./scripts/setup.py --fix        # Fix permissions and ownership
"""

from __future__ import annotations

import argparse
import grp
import os
import pwd
import shutil
import subprocess
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

# Minimum requirements
MIN_PYTHON_VERSION = (3, 10)
REQUIRED_PACKAGES = ["yaml"]  # PyYAML imports as 'yaml'
DOCKER_SERVICE_USER = "docker-services"
DOCKER_GROUP = "docker"
DOCKER_ROOT = Path("/opt/docker")
SYSTEMD_SERVICE_PATH = Path("/etc/systemd/system/docker-compose@.service")
SYSTEMD_INFRA_TARGET_PATH = Path("/etc/systemd/system/docker-compose-infra.target")
CRON_FILE_PATH = Path("/etc/cron.d/docker-maintenance")
LOGROTATE_FILE_PATH = Path("/etc/logrotate.d/ct-controller")

# Infrastructure stacks (start first, before apps)
INFRA_STACKS = ["graylog", "fluentbit", "monitoring", "watchtower"]


@dataclass
class CheckResult:
    """Result of a prerequisite check."""

    name: str
    passed: bool
    message: str
    fixable: bool = False
    fix_command: str | None = None
    details: dict[str, Any] = field(default_factory=dict)


@dataclass
class SetupReport:
    """Complete setup check report."""

    checks: list[CheckResult] = field(default_factory=list)

    @property
    def all_passed(self) -> bool:
        return all(c.passed for c in self.checks)

    @property
    def fixable_issues(self) -> list[CheckResult]:
        return [c for c in self.checks if not c.passed and c.fixable]

    @property
    def blocking_issues(self) -> list[CheckResult]:
        return [c for c in self.checks if not c.passed and not c.fixable]

    def to_dict(self) -> dict:
        return {
            "all_passed": self.all_passed,
            "checks": [
                {
                    "name": c.name,
                    "passed": c.passed,
                    "message": c.message,
                    "fixable": c.fixable,
                    "fix_command": c.fix_command,
                    "details": c.details,
                }
                for c in self.checks
            ],
            "summary": {
                "total": len(self.checks),
                "passed": sum(1 for c in self.checks if c.passed),
                "failed": sum(1 for c in self.checks if not c.passed),
                "fixable": len(self.fixable_issues),
            },
        }


# -----------------------------------------------------------------------------
# Check Functions
# -----------------------------------------------------------------------------


def check_python_version() -> CheckResult:
    """Check Python version meets minimum requirement."""
    current = sys.version_info[:2]
    passed = current >= MIN_PYTHON_VERSION

    return CheckResult(
        name="python_version",
        passed=passed,
        message=(
            f"Python {current[0]}.{current[1]}"
            if passed
            else f"Python {current[0]}.{current[1]} < {MIN_PYTHON_VERSION[0]}.{MIN_PYTHON_VERSION[1]} required"
        ),
        fixable=False,
        details={
            "current": f"{current[0]}.{current[1]}",
            "required": f"{MIN_PYTHON_VERSION[0]}.{MIN_PYTHON_VERSION[1]}",
        },
    )


def check_python_packages() -> CheckResult:
    """Check required Python packages are installed."""
    missing = []
    installed = []

    for package in REQUIRED_PACKAGES:
        try:
            __import__(package)
            installed.append(package)
        except ImportError:
            missing.append(package)

    passed = len(missing) == 0

    # Map import names to pip package names
    pip_names = {"yaml": "PyYAML"}
    pip_missing = [pip_names.get(p, p) for p in missing]

    return CheckResult(
        name="python_packages",
        passed=passed,
        message=(
            f"All packages installed: {', '.join(installed)}"
            if passed
            else f"Missing packages: {', '.join(missing)}"
        ),
        fixable=True,
        fix_command=f"pip install {' '.join(pip_missing)}" if missing else None,
        details={"installed": installed, "missing": missing},
    )


def check_docker_installed() -> CheckResult:
    """Check Docker is installed."""
    docker_path = shutil.which("docker")
    passed = docker_path is not None

    version = None
    if passed:
        try:
            result = subprocess.run(
                ["docker", "--version"],
                capture_output=True,
                text=True,
                timeout=5,
            )
            version = result.stdout.strip()
        except (subprocess.TimeoutExpired, FileNotFoundError):
            pass

    return CheckResult(
        name="docker_installed",
        passed=passed,
        message=version or "Docker installed" if passed else "Docker not found",
        fixable=False,
        details={"path": docker_path, "version": version},
    )


def check_docker_running() -> CheckResult:
    """Check Docker daemon is running."""
    try:
        result = subprocess.run(
            ["docker", "info"],
            capture_output=True,
            text=True,
            timeout=10,
        )
        passed = result.returncode == 0
        message = "Docker daemon running" if passed else "Docker daemon not running"
    except (subprocess.TimeoutExpired, FileNotFoundError):
        passed = False
        message = "Cannot connect to Docker"

    return CheckResult(
        name="docker_running",
        passed=passed,
        message=message,
        fixable=True,
        fix_command="sudo systemctl start docker",
    )


def check_docker_compose() -> CheckResult:
    """Check Docker Compose v2 is available."""
    try:
        result = subprocess.run(
            ["docker", "compose", "version"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        passed = result.returncode == 0
        version = result.stdout.strip() if passed else None
        message = version or "Docker Compose v2 available"
    except (subprocess.TimeoutExpired, FileNotFoundError):
        passed = False
        version = None
        message = "Docker Compose v2 not available"

    return CheckResult(
        name="docker_compose",
        passed=passed,
        message=message,
        fixable=False,
        details={"version": version},
    )


def check_service_user() -> CheckResult:
    """Check docker-services user exists and is in docker group."""
    try:
        user = pwd.getpwnam(DOCKER_SERVICE_USER)
        user_exists = True
        uid = user.pw_uid
        gid = user.pw_gid
    except KeyError:
        user_exists = False
        uid = None
        gid = None

    in_docker_group = False
    if user_exists:
        try:
            docker_group = grp.getgrnam(DOCKER_GROUP)
            in_docker_group = DOCKER_SERVICE_USER in docker_group.gr_mem
        except KeyError:
            pass

    passed = user_exists and in_docker_group
    if not user_exists:
        message = f"User '{DOCKER_SERVICE_USER}' does not exist"
        fix = f"sudo useradd -r -s /usr/sbin/nologin {DOCKER_SERVICE_USER} && sudo usermod -aG docker {DOCKER_SERVICE_USER}"
    elif not in_docker_group:
        message = f"User '{DOCKER_SERVICE_USER}' not in docker group"
        fix = f"sudo usermod -aG docker {DOCKER_SERVICE_USER}"
    else:
        message = f"User '{DOCKER_SERVICE_USER}' configured correctly"
        fix = None

    return CheckResult(
        name="service_user",
        passed=passed,
        message=message,
        fixable=True,
        fix_command=fix,
        details={"user_exists": user_exists, "in_docker_group": in_docker_group, "uid": uid},
    )


def check_docker_root() -> CheckResult:
    """Check /opt/docker directory exists with correct ownership and permissions."""
    exists = DOCKER_ROOT.exists()
    correct_owner = False
    group_writable = False
    has_setgid = False
    owner_info = None
    mode_info = None

    if exists:
        try:
            stat_result = DOCKER_ROOT.stat()
            owner = pwd.getpwuid(stat_result.st_uid).pw_name
            group = grp.getgrgid(stat_result.st_gid).gr_name
            owner_info = f"{owner}:{group}"
            correct_owner = owner == DOCKER_SERVICE_USER and group == DOCKER_SERVICE_USER
            # Check group write (0o020) and setgid (0o2000)
            mode = stat_result.st_mode
            group_writable = bool(mode & 0o020)
            has_setgid = bool(mode & 0o2000)
            mode_info = oct(mode)[-4:]
        except (KeyError, OSError):
            pass

    passed = exists and correct_owner and group_writable and has_setgid
    issues = []

    if not exists:
        message = f"{DOCKER_ROOT} does not exist"
        fix = f"sudo mkdir -p {DOCKER_ROOT} && sudo chown {DOCKER_SERVICE_USER}:{DOCKER_SERVICE_USER} {DOCKER_ROOT} && sudo chmod 2775 {DOCKER_ROOT}"
    else:
        if not correct_owner:
            issues.append(f"ownership is {owner_info}")
        if not group_writable:
            issues.append("missing group write")
        if not has_setgid:
            issues.append("missing setgid")

        if issues:
            message = f"{DOCKER_ROOT}: {', '.join(issues)}"
            # Comprehensive fix: ownership, group write, and setgid recursively
            fix = (
                f"sudo chown -R {DOCKER_SERVICE_USER}:{DOCKER_SERVICE_USER} {DOCKER_ROOT} && "
                f"sudo chmod -R g+w {DOCKER_ROOT} && "
                f"sudo find {DOCKER_ROOT} -type d -exec chmod g+s {{}} \\;"
            )
        else:
            message = f"{DOCKER_ROOT} exists with correct ownership and permissions"
            fix = None

    return CheckResult(
        name="docker_root",
        passed=passed,
        message=message,
        fixable=True,
        fix_command=fix,
        details={
            "exists": exists,
            "owner": owner_info,
            "mode": mode_info,
            "group_writable": group_writable,
            "setgid": has_setgid,
        },
    )


def check_systemd_service() -> CheckResult:
    """Check systemd service template is installed."""
    installed = SYSTEMD_SERVICE_PATH.exists()

    # Check if source exists in this repo (now in scripts/templates/)
    script_dir = Path(__file__).parent
    source_path = script_dir / "templates" / "docker-compose@.service"
    source_exists = source_path.exists()

    passed = installed
    if not installed and source_exists:
        message = "Systemd service template not installed"
        fix = f"sudo cp {source_path} {SYSTEMD_SERVICE_PATH} && sudo systemctl daemon-reload"
    elif not installed:
        message = "Systemd service template not installed (source not found)"
        fix = None
    else:
        message = "Systemd service template installed"
        fix = None

    return CheckResult(
        name="systemd_service",
        passed=passed,
        message=message,
        fixable=source_exists,
        fix_command=fix,
        details={"installed": installed, "path": str(SYSTEMD_SERVICE_PATH)},
    )


def check_cron_installed() -> CheckResult:
    """Check maintenance cron jobs are installed."""
    installed = CRON_FILE_PATH.exists()

    # Check if source exists
    script_dir = Path(__file__).parent
    source_path = script_dir / "templates" / "docker-maintenance.cron"
    source_exists = source_path.exists()

    passed = installed
    if not installed and source_exists:
        message = "Maintenance cron jobs not installed"
        fix = f"sudo cp {source_path} {CRON_FILE_PATH} && sudo chmod 644 {CRON_FILE_PATH}"
    elif not installed:
        message = "Maintenance cron jobs not installed (source not found)"
        fix = None
    else:
        message = "Maintenance cron jobs installed"
        fix = None

    return CheckResult(
        name="cron_jobs",
        passed=passed,
        message=message,
        fixable=source_exists,
        fix_command=fix,
        details={"installed": installed, "path": str(CRON_FILE_PATH)},
    )


def check_logrotate_installed() -> CheckResult:
    """Check logrotate config is installed."""
    installed = LOGROTATE_FILE_PATH.exists()

    # Check if source exists
    script_dir = Path(__file__).parent
    source_path = script_dir / "templates" / "ct-controller-logrotate.conf"
    source_exists = source_path.exists()

    passed = installed
    if not installed and source_exists:
        message = "Logrotate config not installed"
        fix = f"sudo cp {source_path} {LOGROTATE_FILE_PATH} && sudo chmod 644 {LOGROTATE_FILE_PATH}"
    elif not installed:
        message = "Logrotate config not installed (source not found)"
        fix = None
    else:
        message = "Logrotate config installed"
        fix = None

    return CheckResult(
        name="logrotate",
        passed=passed,
        message=message,
        fixable=source_exists,
        fix_command=fix,
        details={"installed": installed, "path": str(LOGROTATE_FILE_PATH)},
    )


def check_current_user_docker() -> CheckResult:
    """Check if current user can run Docker commands."""
    try:
        result = subprocess.run(
            ["docker", "ps"],
            capture_output=True,
            text=True,
            timeout=10,
        )
        passed = result.returncode == 0
        if passed:
            message = "Current user can run Docker commands"
        else:
            message = "Current user cannot run Docker (permission denied?)"
    except (subprocess.TimeoutExpired, FileNotFoundError):
        passed = False
        message = "Cannot run Docker commands"

    current_user = os.getenv("USER", "unknown")

    return CheckResult(
        name="current_user_docker",
        passed=passed,
        message=message,
        fixable=True,
        fix_command=f"sudo usermod -aG docker {current_user} && newgrp docker",
        details={"user": current_user},
    )


def check_boot_ordering() -> CheckResult:
    """Check if boot ordering (infra target + app overrides) is installed."""
    script_dir = Path(__file__).parent

    # Check if infra target exists
    infra_target_installed = SYSTEMD_INFRA_TARGET_PATH.exists()
    infra_target_source = script_dir / "templates" / "docker-compose-infra.target"

    # Find app stacks that need overrides (non-infra stacks with docker-compose.yml)
    # Includes nested stacks like stack/substack
    app_stacks = []
    if DOCKER_ROOT.exists():
        for stack_dir in DOCKER_ROOT.iterdir():
            if stack_dir.is_dir() and not stack_dir.name.startswith("."):
                # Skip non-stack directories
                if stack_dir.name in ("scripts", "tests", "docs", "node_modules"):
                    continue

                compose_file = stack_dir / "docker-compose.yml"
                if compose_file.exists() and stack_dir.name not in INFRA_STACKS:
                    app_stacks.append(stack_dir.name)

                # Check for nested stacks (e.g., stack/substack)
                for substack_dir in stack_dir.iterdir():
                    if substack_dir.is_dir() and not substack_dir.name.startswith("."):
                        substack_compose = substack_dir / "docker-compose.yml"
                        if substack_compose.exists():
                            stack_name = f"{stack_dir.name}-{substack_dir.name}"
                            if stack_name not in app_stacks:
                                app_stacks.append(stack_name)

    # Check which app stacks have overrides installed
    missing_overrides = []
    for stack in app_stacks:
        override_path = Path(f"/etc/systemd/system/docker-compose@{stack}.service.d/after-infra.conf")
        if not override_path.exists():
            missing_overrides.append(stack)

    # Determine status
    passed = infra_target_installed and len(missing_overrides) == 0

    if not infra_target_installed:
        message = "Boot ordering not installed (infra target missing)"
        fix = "install_boot_ordering"  # Special marker for custom fix function
    elif missing_overrides:
        message = f"Boot ordering incomplete ({len(missing_overrides)} app stacks missing overrides)"
        fix = "install_boot_ordering"
    else:
        message = "Boot ordering installed (infra target + app overrides)"
        fix = None

    return CheckResult(
        name="boot_ordering",
        passed=passed,
        message=message,
        fixable=infra_target_source.exists(),
        fix_command=fix,
        details={
            "infra_target_installed": infra_target_installed,
            "app_stacks": app_stacks,
            "missing_overrides": missing_overrides,
        },
    )


# -----------------------------------------------------------------------------
# Setup Actions
# -----------------------------------------------------------------------------


def run_all_checks() -> SetupReport:
    """Run all prerequisite checks."""
    report = SetupReport()

    report.checks.append(check_python_version())
    report.checks.append(check_python_packages())
    report.checks.append(check_docker_installed())
    report.checks.append(check_docker_running())
    report.checks.append(check_docker_compose())
    report.checks.append(check_current_user_docker())
    report.checks.append(check_service_user())
    report.checks.append(check_docker_root())
    report.checks.append(check_systemd_service())
    report.checks.append(check_cron_installed())
    report.checks.append(check_logrotate_installed())
    report.checks.append(check_boot_ordering())

    return report


def install_boot_ordering() -> bool:
    """Install boot ordering (infra target + app overrides)."""
    script_dir = Path(__file__).parent
    infra_target_source = script_dir / "templates" / "docker-compose-infra.target"
    override_source = script_dir / "templates" / "docker-compose-app-override.conf"

    try:
        # Install infra target
        if not SYSTEMD_INFRA_TARGET_PATH.exists():
            result = subprocess.run(
                ["sudo", "cp", str(infra_target_source), str(SYSTEMD_INFRA_TARGET_PATH)],
                capture_output=True,
                text=True,
            )
            if result.returncode != 0:
                print(f"  Failed to copy infra target: {result.stderr}")
                return False

        # Find app stacks (including nested stacks)
        app_stacks = []
        if DOCKER_ROOT.exists():
            for stack_dir in DOCKER_ROOT.iterdir():
                if stack_dir.is_dir() and not stack_dir.name.startswith("."):
                    # Skip non-stack directories
                    if stack_dir.name in ("scripts", "tests", "docs", "node_modules"):
                        continue

                    compose_file = stack_dir / "docker-compose.yml"
                    if compose_file.exists() and stack_dir.name not in INFRA_STACKS:
                        app_stacks.append(stack_dir.name)

                    # Check for nested stacks (e.g., stack/substack)
                    for substack_dir in stack_dir.iterdir():
                        if substack_dir.is_dir() and not substack_dir.name.startswith("."):
                            substack_compose = substack_dir / "docker-compose.yml"
                            if substack_compose.exists():
                                stack_name = f"{stack_dir.name}-{substack_dir.name}"
                                if stack_name not in app_stacks:
                                    app_stacks.append(stack_name)

        # Install app overrides
        for stack in app_stacks:
            override_dir = Path(f"/etc/systemd/system/docker-compose@{stack}.service.d")
            override_path = override_dir / "after-infra.conf"

            if not override_path.exists():
                # Create directory
                subprocess.run(["sudo", "mkdir", "-p", str(override_dir)], capture_output=True)

                # Copy override
                result = subprocess.run(
                    ["sudo", "cp", str(override_source), str(override_path)],
                    capture_output=True,
                    text=True,
                )
                if result.returncode != 0:
                    print(f"  Failed to install override for {stack}: {result.stderr}")
                    return False

        # Reload systemd
        result = subprocess.run(
            ["sudo", "systemctl", "daemon-reload"],
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            print(f"  Failed to reload systemd: {result.stderr}")
            return False

        return True

    except Exception as e:
        print(f"  Error: {e}")
        return False


def run_fix(check: CheckResult) -> bool:
    """Attempt to fix a failed check."""
    if not check.fix_command:
        return False

    # Special handling for boot ordering
    if check.fix_command == "install_boot_ordering":
        print(f"  Installing boot ordering (infra target + app overrides)")
        return install_boot_ordering()

    print(f"  Running: {check.fix_command}")
    try:
        result = subprocess.run(
            check.fix_command,
            shell=True,
            capture_output=True,
            text=True,
        )
        return result.returncode == 0
    except Exception as e:
        print(f"  Error: {e}")
        return False


# -----------------------------------------------------------------------------
# Output
# -----------------------------------------------------------------------------


def report_human(report: SetupReport, verbose: bool = False) -> None:
    """Print human-readable report."""
    GREEN = "\033[0;32m"
    RED = "\033[0;31m"
    YELLOW = "\033[1;33m"
    CYAN = "\033[0;36m"
    BOLD = "\033[1m"
    RESET = "\033[0m"

    print()
    print(f"{BOLD}ct-controller Setup Check{RESET}")
    print("=" * 50)
    print()

    for check in report.checks:
        if check.passed:
            status = f"{GREEN}[OK]{RESET}"
        elif check.fixable:
            status = f"{YELLOW}[FIX]{RESET}"
        else:
            status = f"{RED}[FAIL]{RESET}"

        print(f"  {status} {check.name}: {check.message}")

        if verbose and check.details:
            for key, value in check.details.items():
                print(f"       {key}: {value}")

        if not check.passed and check.fix_command:
            print(f"       {CYAN}Fix: {check.fix_command}{RESET}")

    print()
    print("=" * 50)

    summary = report.to_dict()["summary"]
    print(f"  Total: {summary['total']}, Passed: {summary['passed']}, Failed: {summary['failed']}")

    if report.all_passed:
        print(f"\n{GREEN}All checks passed - system ready{RESET}")
    elif report.fixable_issues:
        print(f"\n{YELLOW}{len(report.fixable_issues)} issue(s) can be fixed automatically{RESET}")
        print(f"Run with {BOLD}--install{RESET} to fix")
    else:
        print(f"\n{RED}Some issues require manual intervention{RESET}")


def report_json(report: SetupReport) -> None:
    """Print JSON report."""
    import json

    print(json.dumps(report.to_dict(), indent=2))


# -----------------------------------------------------------------------------
# Main
# -----------------------------------------------------------------------------


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Setup and prerequisites check for ct-controller",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s                  Check prerequisites (no changes)
  %(prog)s --install        Install missing packages and configure
  %(prog)s --fix            Fix permissions only
  %(prog)s --json           Output JSON report
""",
    )
    parser.add_argument(
        "--install",
        action="store_true",
        help="Install missing dependencies and configure system",
    )
    parser.add_argument(
        "--fix",
        action="store_true",
        help="Fix permissions and ownership issues only",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Output JSON format",
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Show detailed information",
    )

    args = parser.parse_args()

    # Run checks
    report = run_all_checks()

    # Handle install/fix modes
    if args.install or args.fix:
        if not args.json:
            print("\nAttempting to fix issues...\n")

        fixed = 0
        for check in report.checks:
            if not check.passed and check.fixable and check.fix_command:
                if args.fix and check.name not in ("docker_root", "service_user", "systemd_service"):
                    continue  # --fix only does permission-related fixes

                if not args.json:
                    print(f"Fixing: {check.name}")

                if run_fix(check):
                    fixed += 1
                    if not args.json:
                        print(f"  Success")
                else:
                    if not args.json:
                        print(f"  Failed (may need sudo)")

        if not args.json:
            print(f"\nFixed {fixed} issue(s)")
            print("Re-running checks...\n")

        # Re-run checks
        report = run_all_checks()

    # Output
    if args.json:
        report_json(report)
    else:
        report_human(report, verbose=args.verbose)

    return 0 if report.all_passed else 1


if __name__ == "__main__":
    sys.exit(main())