ArdupilotCustomFirmwareBuilder/metadata_manager/firmware_server/client.py

import json
import logging
import os
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import requests

from .exceptions import ManifestFetchError


@dataclass
class _CacheMeta:
    etag: Optional[str] = None
    last_modified: Optional[str] = None
    fetched_at: Optional[str] = None

    def to_dict(self) -> dict:
        return {
            "etag": self.etag,
            "last_modified": self.last_modified,
            "fetched_at": self.fetched_at,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "_CacheMeta":
        return cls(
            etag=data.get("etag"),
            last_modified=data.get("last_modified"),
            fetched_at=data.get("fetched_at"),
        )


class ManifestClient:
    """Fetch and cache the ArduPilot firmware manifest.json file."""

    def __init__(
        self,
        url: str,
        cache_dir: str,
        timeout: int = 120,
        user_agent: str = "CustomBuild/1.0",
    ):
        self.url = url
        self.cache_dir = Path(cache_dir)
        self.cache_path = self.cache_dir / "manifest.json"
        self.meta_path = self.cache_dir / "manifest.json.meta"
        self.timeout = timeout
        self.user_agent = user_agent
        self.logger = logging.getLogger(__name__)

    def fetch_raw(self) -> bytes:
        headers = {"User-Agent": self.user_agent}
        meta = self._read_meta() if self._has_cache() else _CacheMeta()

        if meta.etag:
            headers["If-None-Match"] = meta.etag
        if meta.last_modified:
            headers["If-Modified-Since"] = meta.last_modified

        try:
            response = requests.get(
                self.url,
                headers=headers,
                timeout=self.timeout,
            )
        except requests.RequestException as exc:
            return self._fallback_or_raise(exc)

        if response.status_code == 304:
            self.logger.info("Manifest not modified (304), using cache")
            self._touch_fetched_at(self._now_iso())
            return self._read_cache_bytes()

        if response.status_code != 200:
            return self._fallback_or_raise(
                ManifestFetchError(
                    f"Manifest fetch failed with status {response.status_code}"
                )
            )

        raw = response.content
        self._write_cache(
            raw,
            _CacheMeta(
                etag=response.headers.get("ETag"),
                last_modified=response.headers.get("Last-Modified"),
                fetched_at=self._now_iso(),
            ),
        )
        self.logger.info("Downloaded manifest (%d bytes)", len(raw))
        return raw

    def fetch(self) -> dict:
        return json.loads(self.fetch_raw().decode("utf-8"))

    def _has_cache(self) -> bool:
        return self.cache_path.is_file()

    def _read_cache_bytes(self) -> bytes:
        return self.cache_path.read_bytes()

    def _read_meta(self) -> _CacheMeta:
        if not self.meta_path.is_file():
            return _CacheMeta()
        return _CacheMeta.from_dict(
            json.loads(self.meta_path.read_text(encoding="utf-8"))
        )

    def _write_cache(self, raw: bytes, meta: _CacheMeta) -> None:
        self._atomic_write_bytes(self.cache_path, raw)
        self._atomic_write_text(
            self.meta_path,
            json.dumps(meta.to_dict(), indent=2),
        )

    def _touch_fetched_at(self, fetched_at: str) -> None:
        meta = self._read_meta()
        meta.fetched_at = fetched_at
        self._atomic_write_text(
            self.meta_path,
            json.dumps(meta.to_dict(), indent=2),
        )

    def _atomic_write_bytes(self, path: Path, raw: bytes) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp_path = path.with_name(f"{path.name}.tmp")
        tmp_path.write_bytes(raw)
        os.replace(tmp_path, path)

    def _atomic_write_text(self, path: Path, text: str) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp_path = path.with_name(f"{path.name}.tmp")
        tmp_path.write_text(text, encoding="utf-8")
        os.replace(tmp_path, path)

    def _fallback_or_raise(self, exc: Exception) -> bytes:
        if self._has_cache():
            self.logger.warning(
                "Manifest fetch failed (%s), using stale cache", exc
            )
            return self._read_cache_bytes()
        raise ManifestFetchError(str(exc)) from exc

    @staticmethod
    def _now_iso() -> str:
        return datetime.now(timezone.utc).isoformat()