kiddo/backend/update.py

import json
import os
import subprocess
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List

import httpx

from backend.settings import Settings


def _project_root() -> Path:
    return Path(__file__).resolve().parents[1]


def _read_version() -> str:
    try:
        return (_project_root() / "VERSION").read_text(encoding="utf-8").strip()
    except OSError:
        return "unknown"


def _status_path(settings: Settings) -> Path:
    return Path(settings.update_status_file)


def _log_path(settings: Settings) -> Path:
    return Path(settings.update_log_file)


def _ensure_parent(path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)


def _write_status(settings: Settings, status: str, version: str, error: str | None = None) -> None:
    status_path = _status_path(settings)
    _ensure_parent(status_path)
    payload = {
        "device_id": os.uname().nodename,
        "version": version,
        "status": status,
        "error": error or "",
        "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
    }
    status_path.write_text(json.dumps(payload), encoding="utf-8")
    _append_log(settings, payload)


def _append_log(settings: Settings, payload: Dict[str, Any]) -> None:
    log_path = _log_path(settings)
    _ensure_parent(log_path)
    with log_path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(payload) + "\n")


def get_status(settings: Settings) -> Dict[str, Any]:
    current_version = _read_version()
    status_path = _status_path(settings)
    if status_path.exists():
        try:
            data = json.loads(status_path.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            data = {}
    else:
        data = {}

    return {
        "current_version": current_version,
        "last_status": data.get("status", "unknown"),
        "last_error": data.get("error") or None,
        "last_timestamp": data.get("timestamp"),
        "enrolled": bool(settings.update_token),
    }


def enroll(settings: Settings) -> str:
    if not settings.update_enroll_token:
        raise ValueError("No enrollment token provided in settings")

    enroll_url = f"{settings.update_service_url}/v1/enroll"
    payload = {
        "project_id": settings.update_project_id,
        "client_id": os.uname().nodename,
        "software_id": "safe-kiddo",
        "enroll_token": settings.update_enroll_token,
    }

    with httpx.Client(timeout=10.0) as client:
        response = client.post(enroll_url, json=payload)
        response.raise_for_status()
        data = response.json()

    token = data.get("token")
    if not token:
        raise ValueError("Enrollment response did not contain a token")

    # Save token
    token_path = Path(settings.update_token_file)
    _ensure_parent(token_path)
    token_path.write_text(token, encoding="utf-8")
    # Update settings object for immediate use
    settings.update_token = token

    return token


def _parse_version(value: str) -> List[int]:
    return [int(part) for part in value.split(".")]


def _get_fresh_token(settings: Settings) -> str:
    """Always reload token from file/env to avoid stale cache."""
    # Env var takes precedence
    env_token = os.getenv("SKD_UPDATE_TOKEN", "")
    if env_token:
        return env_token
    # Fallback to file
    token_file = Path(settings.update_token_file)
    if token_file.exists():
        try:
            return token_file.read_text(encoding="utf-8").strip()
        except OSError:
            pass
    return settings.update_token  # Fallback to cached value


def check_update(settings: Settings) -> Dict[str, Any]:
    token = _get_fresh_token(settings)
    if not token:
        raise ValueError("Client is not enrolled (missing update token)")

    headers = {"Authorization": f"Bearer {token}"}
    manifest_url = (
        f"{settings.update_service_url}/v1/projects/{settings.update_project_id}/manifest"
    )

    with httpx.Client(timeout=10.0) as client:
        response = client.get(manifest_url, headers=headers)
        response.raise_for_status()
        manifest = response.json()

    latest_version = manifest.get("version", "")
    artifact_url = manifest.get("artifact_url", "")
    sha256 = manifest.get("sha256", "")
    message = manifest.get("message")

    available = False
    current_version = _read_version()
    try:
        available = _parse_version(latest_version) > _parse_version(current_version)
    except ValueError:
        if latest_version and latest_version != current_version:
            available = True

    return {
        "available": available,
        "latest_version": latest_version,
        "artifact_url": artifact_url,
        "sha256": sha256,
        "message": message,
    }


def report_status(
    settings: Settings,
    status: str,
    version: str,
    error: str | None = None,
    duration_ms: int | None = None,
) -> None:
    token = _get_fresh_token(settings)
    if not token:
        return

    report_url = (
        f"{settings.update_service_url}/v1/projects/{settings.update_project_id}/status"
    )
    payload = {
        "project_id": settings.update_project_id,
        "version": version,
        "status": status,
        "timestamp": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "client_id": os.uname().nodename,
    }
    if error:
        payload["error_code"] = error
        payload["reason"] = error
    if duration_ms is not None:
        payload["duration_ms"] = duration_ms

    try:
        headers = {"Authorization": f"Bearer {token}"}
        with httpx.Client(timeout=10.0) as client:
            client.post(report_url, json=payload, headers=headers).raise_for_status()
    except Exception:
        # We don't want to crash if status reporting fails
        pass


def _run_async(script_path: Path, settings: Settings) -> None:
    # Explicitly gather the env vars we need to pass
    env_vars = {
        "SKD_UPDATE_SERVICE_URL": settings.update_service_url,
        "SKD_UPDATE_PROJECT_ID": settings.update_project_id,
        "SKD_UPDATE_TOKEN": _get_fresh_token(settings),
        "SKD_UPDATE_STATUS_FILE": settings.update_status_file,
        "SKD_UPDATE_LOG_FILE": settings.update_log_file,
    }

    # Construct systemd-run command with --setenv for each variable
    cmd = [
        "systemd-run",
        "--unit=skd-update",
        "--collect",
        "--description=Safe Kiddo Update Process",
    ]
    
    for key, val in env_vars.items():
        if val:  # Only pass if not empty
            cmd.append(f"--setenv={key}={val}")
            
    cmd.append(str(script_path))

    # We don't pass 'env' parameter to Popen because systemd-run ignores it 
    # for the target process (it only uses it for itself, but we use --setenv).
    subprocess.Popen(
        cmd,
        cwd="/",
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )


def start_update(settings: Settings, version: str | None = None) -> None:
    current_version = version or _read_version()
    _write_status(settings, "in_progress", current_version)
    report_status(settings, "in_progress", current_version)
    script = _project_root() / "scripts" / "update_client.sh"
    _run_async(script, settings)


def start_rollback(settings: Settings) -> None:
    current_version = _read_version()
    _write_status(settings, "in_progress", current_version)
    report_status(settings, "in_progress", current_version)
    script = _project_root() / "scripts" / "rollback_client.sh"
    _run_async(script, settings)


def get_logs(settings: Settings, limit: int = 200) -> List[Dict[str, Any]]:
    log_path = _log_path(settings)
    if not log_path.exists():
        return []
    lines = log_path.read_text(encoding="utf-8").splitlines()
    entries: List[Dict[str, Any]] = []
    for line in lines[-limit:]:
        try:
            entries.append(json.loads(line))
        except json.JSONDecodeError:
            continue
    return entries


def get_service_status(settings: Settings) -> Dict[str, Any]:
    base_url = settings.update_service_url.rstrip("/")
    if not base_url:
        return {
            "url": "",
            "reachable": False,
            "status_code": None,
            "error": "update service url not configured",
            "checked_url": "",
            "environment": "unknown",
        }

    env = "prod"
    lowered = base_url.lower()
    if "://dev." in lowered or lowered.startswith("dev."):
        env = "dev"
    elif "://staging." in lowered or lowered.startswith("staging."):
        env = "staging"

    check_url = base_url
    try:
        with httpx.Client(timeout=3.0) as client:
            response = client.get(check_url)
        return {
            "url": base_url,
            "reachable": True,
            "status_code": response.status_code,
            "error": None,
            "checked_url": check_url,
            "environment": env,
        }
    except Exception as exc:
        return {
            "url": base_url,
            "reachable": False,
            "status_code": None,
            "error": str(exc),
            "checked_url": check_url,
            "environment": env,
        }