unifi-network-maps 1.4.1__py3-none-any.whl → 1.4.2__py3-none-any.whl

unifi_network_maps/__init__.py

@@ -1 +1 @@
-__version__ = "1.4.1"
+__version__ = "1.4.2"

unifi_network_maps/adapters/unifi.py

@@ -3,13 +3,17 @@
 from __future__ import annotations
 
 import hashlib
+import json
 import logging
 import os
-import pickle
+import stat
 import time
-from collections.abc import Iterable
+from collections.abc import Callable, Iterator, Sequence
+from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import TimeoutError as FutureTimeoutError
+from contextlib import contextmanager
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import IO, TYPE_CHECKING
 
 from .config import Config
 
@@ -23,6 +27,194 @@ def _cache_dir() -> Path:
     return Path(os.environ.get("UNIFI_CACHE_DIR", ".cache/unifi_network_maps"))
 
 
+def _device_attr(device: object, name: str) -> object | None:
+    if isinstance(device, dict):
+        return device.get(name)
+    return getattr(device, name, None)
+
+
+def _first_attr(device: object, *names: str) -> object | None:
+    for name in names:
+        value = _device_attr(device, name)
+        if value is not None:
+            return value
+    return None
+
+
+def _as_list(value: object | None) -> list[object]:
+    if value is None:
+        return []
+    if isinstance(value, list):
+        return value
+    if isinstance(value, dict):
+        return [value]
+    if isinstance(value, str | bytes):
+        return []
+    try:
+        return list(value)  # type: ignore[arg-type]
+    except TypeError:
+        return []
+
+
+def _serialize_lldp_entry(entry: object) -> dict[str, object]:
+    return {
+        "chassis_id": _first_attr(entry, "chassis_id", "chassisId"),
+        "port_id": _first_attr(entry, "port_id", "portId"),
+        "port_desc": _first_attr(entry, "port_desc", "portDesc", "port_descr", "portDescr"),
+        "local_port_name": _first_attr(entry, "local_port_name", "localPortName"),
+        "local_port_idx": _first_attr(entry, "local_port_idx", "localPortIdx"),
+    }
+
+
+def _serialize_lldp_entries(value: object | None) -> list[dict[str, object]]:
+    entries = _as_list(value)
+    serialized: list[dict[str, object]] = []
+    for entry in entries:
+        data = _serialize_lldp_entry(entry)
+        if data.get("chassis_id") and data.get("port_id"):
+            serialized.append(data)
+    return serialized
+
+
+def _serialize_port_entry(entry: object) -> dict[str, object]:
+    aggregation_group = _first_attr(
+        entry,
+        "aggregation_group",
+        "aggregation_id",
+        "aggregate_id",
+        "agg_id",
+        "lag_id",
+        "lag_group",
+        "link_aggregation_group",
+        "link_aggregation_id",
+        "aggregate",
+        "aggregated_by",
+    )
+    return {
+        "port_idx": _first_attr(entry, "port_idx", "portIdx"),
+        "name": _first_attr(entry, "name"),
+        "ifname": _first_attr(entry, "ifname"),
+        "speed": _first_attr(entry, "speed"),
+        "aggregation_group": aggregation_group,
+        "port_poe": _first_attr(entry, "port_poe"),
+        "poe_enable": _first_attr(entry, "poe_enable"),
+        "poe_good": _first_attr(entry, "poe_good"),
+        "poe_power": _first_attr(entry, "poe_power"),
+    }
+
+
+def _serialize_port_table(value: object | None) -> list[dict[str, object]]:
+    return [_serialize_port_entry(entry) for entry in _as_list(value)]
+
+
+def _serialize_uplink(value: object | None) -> dict[str, object] | None:
+    if value is None:
+        return None
+    data = {
+        "uplink_mac": _first_attr(value, "uplink_mac", "uplink_device_mac"),
+        "uplink_device_name": _first_attr(value, "uplink_device_name", "uplink_name"),
+        "uplink_remote_port": _first_attr(value, "uplink_remote_port", "port_idx"),
+    }
+    if any(item is not None for item in data.values()):
+        return data
+    return None
+
+
+def _device_lldp_value(device: object) -> object | None:
+    lldp_info = _device_attr(device, "lldp_info")
+    if lldp_info is None:
+        lldp_info = _device_attr(device, "lldp")
+    if lldp_info is None:
+        lldp_info = _device_attr(device, "lldp_table")
+    return lldp_info
+
+
+def _device_uplink_fields(device: object) -> dict[str, object | None]:
+    return {
+        "uplink": _serialize_uplink(_device_attr(device, "uplink")),
+        "last_uplink": _serialize_uplink(_device_attr(device, "last_uplink")),
+        "uplink_mac": _first_attr(device, "uplink_mac", "uplink_device_mac"),
+        "uplink_device_name": _device_attr(device, "uplink_device_name"),
+        "uplink_remote_port": _device_attr(device, "uplink_remote_port"),
+        "last_uplink_mac": _device_attr(device, "last_uplink_mac"),
+    }
+
+
+def _serialize_device_for_cache(device: object) -> dict[str, object]:
+    payload = {
+        "name": _device_attr(device, "name"),
+        "model_name": _device_attr(device, "model_name"),
+        "model": _device_attr(device, "model"),
+        "mac": _device_attr(device, "mac"),
+        "ip": _first_attr(device, "ip", "ip_address"),
+        "type": _first_attr(device, "type", "device_type"),
+        "displayable_version": _first_attr(device, "displayable_version", "version"),
+        "lldp_info": _serialize_lldp_entries(_device_lldp_value(device)),
+        "port_table": _serialize_port_table(_device_attr(device, "port_table")),
+    }
+    payload.update(_device_uplink_fields(device))
+    return payload
+
+
+def _serialize_devices_for_cache(devices: Sequence[object]) -> list[dict[str, object]]:
+    return [_serialize_device_for_cache(device) for device in devices]
+
+
+def _cache_lock_path(path: Path) -> Path:
+    return path.with_suffix(path.suffix + ".lock")
+
+
+def _acquire_cache_lock(lock_file: IO[str]) -> None:
+    if os.name == "nt":
+        import msvcrt
+
+        msvcrt.locking(lock_file.fileno(), msvcrt.LK_LOCK, 1)
+    else:
+        import fcntl
+
+        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
+
+
+def _release_cache_lock(lock_file: IO[str]) -> None:
+    if os.name == "nt":
+        import msvcrt
+
+        msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1)
+    else:
+        import fcntl
+
+        fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
+
+
+@contextmanager
+def _cache_lock(path: Path) -> Iterator[None]:
+    lock_path = _cache_lock_path(path)
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+    with lock_path.open("a+", encoding="utf-8") as lock_file:
+        try:
+            _acquire_cache_lock(lock_file)
+            yield
+        finally:
+            try:
+                _release_cache_lock(lock_file)
+            except OSError:
+                logger.debug("Failed to release cache lock %s", lock_path)
+
+
+def _is_cache_dir_safe(path: Path) -> bool:
+    if not path.exists():
+        return True
+    try:
+        mode = stat.S_IMODE(path.stat().st_mode)
+    except OSError as exc:
+        logger.warning("Failed to stat cache dir %s: %s", path, exc)
+        return False
+    if mode & (stat.S_IWGRP | stat.S_IWOTH):
+        logger.warning("Cache dir %s is group/world-writable; skipping cache", path)
+        return False
+    return True
+
+
 def _cache_ttl_seconds() -> int:
     value = os.environ.get("UNIFI_CACHE_TTL_SECONDS", "").strip()
     if not value:
@@ -38,7 +230,7 @@ def _cache_key(*parts: str) -> str:
     return digest[:24]
 
 
-def _load_cache(path: Path, ttl_seconds: int) -> object | None:
+def _load_cache(path: Path, ttl_seconds: int) -> Sequence[object] | None:
     data, age = _load_cache_with_age(path)
     if data is None:
         return None
@@ -49,14 +241,18 @@ def _load_cache(path: Path, ttl_seconds: int) -> object | None:
     return data
 
 
-def _load_cache_with_age(path: Path) -> tuple[object | None, float | None]:
+def _load_cache_with_age(path: Path) -> tuple[Sequence[object] | None, float | None]:
     if not path.exists():
         return None, None
     try:
-        payload = pickle.loads(path.read_bytes())
+        with _cache_lock(path):
+            payload = json.loads(path.read_text(encoding="utf-8"))
     except Exception as exc:
         logger.debug("Failed to read cache %s: %s", path, exc)
         return None, None
+    if not isinstance(payload, dict):
+        logger.debug("Cached payload at %s is not a dict", path)
+        return None, None
     timestamp = payload.get("timestamp")
     if not isinstance(timestamp, int | float):
         return None, None
@@ -67,13 +263,16 @@ def _load_cache_with_age(path: Path) -> tuple[object | None, float | None]:
     return data, time.time() - timestamp
 
 
-def _save_cache(path: Path, data: object) -> None:
+def _save_cache(path: Path, data: Sequence[object]) -> None:
     try:
         path.parent.mkdir(parents=True, exist_ok=True)
+        if not _is_cache_dir_safe(path.parent):
+            return
         payload = {"timestamp": time.time(), "data": data}
         tmp_path = path.with_suffix(path.suffix + ".tmp")
-        tmp_path.write_bytes(pickle.dumps(payload))
-        tmp_path.replace(path)
+        with _cache_lock(path):
+            tmp_path.write_text(json.dumps(payload, ensure_ascii=True), encoding="utf-8")
+            tmp_path.replace(path)
     except Exception as exc:
         logger.debug("Failed to write cache %s: %s", path, exc)
 
@@ -99,13 +298,37 @@ def _retry_backoff_seconds() -> float:
         return 0.5
 
 
-def _call_with_retries(operation: str, func) -> object:
+def _request_timeout_seconds() -> float | None:
+    value = os.environ.get("UNIFI_REQUEST_TIMEOUT_SECONDS", "").strip()
+    if not value:
+        return None
+    try:
+        return max(0.0, float(value))
+    except ValueError:
+        logger.warning("Invalid UNIFI_REQUEST_TIMEOUT_SECONDS value: %s", value)
+        return None
+
+
+def _call_with_timeout[T](operation: str, func: Callable[[], T], timeout: float | None) -> T:
+    if timeout is None or timeout <= 0:
+        return func()
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        future = executor.submit(func)
+        try:
+            return future.result(timeout=timeout)
+        except FutureTimeoutError as exc:
+            future.cancel()
+            raise TimeoutError(f"{operation} timed out after {timeout:.2f}s") from exc
+
+
+def _call_with_retries[T](operation: str, func: Callable[[], T]) -> T:
     attempts = _retry_attempts()
     backoff = _retry_backoff_seconds()
+    timeout = _request_timeout_seconds()
     last_exc: Exception | None = None
     for attempt in range(1, attempts + 1):
         try:
-            return func()
+            return _call_with_timeout(operation, func, timeout)
         except Exception as exc:  # noqa: BLE001 - surface full error after retries
             last_exc = exc
             logger.warning("Failed %s attempt %d/%d: %s", operation, attempt, attempts, exc)
@@ -129,8 +352,12 @@ def _init_controller(config: Config, *, is_udm_pro: bool) -> UnifiController:
 
 
 def fetch_devices(
-    config: Config, *, site: str | None = None, detailed: bool = True
-) -> Iterable[object]:
+    config: Config,
+    *,
+    site: str | None = None,
+    detailed: bool = True,
+    use_cache: bool = True,
+) -> Sequence[object]:
     """Fetch devices from UniFi Controller.
 
     Uses `unifi-controller-api` to authenticate and return device objects.
@@ -142,9 +369,13 @@ def fetch_devices(
 
     site_name = site or config.site
     ttl_seconds = _cache_ttl_seconds()
-    cache_path = _cache_dir() / f"devices_{_cache_key(config.url, site_name, str(detailed))}.pkl"
-    cached = _load_cache(cache_path, ttl_seconds)
-    stale_cached, cache_age = _load_cache_with_age(cache_path)
+    cache_path = _cache_dir() / f"devices_{_cache_key(config.url, site_name, str(detailed))}.json"
+    if use_cache and _is_cache_dir_safe(cache_path.parent):
+        cached = _load_cache(cache_path, ttl_seconds)
+        stale_cached, cache_age = _load_cache_with_age(cache_path)
+    else:
+        cached = None
+        stale_cached, cache_age = None, None
     if cached is not None:
         logger.info("Using cached devices (%d)", len(cached))
         return cached
@@ -155,7 +386,7 @@ def fetch_devices(
         logger.info("UDM Pro authentication failed, retrying legacy auth")
         controller = _init_controller(config, is_udm_pro=False)
 
-    def _fetch() -> list[object]:
+    def _fetch() -> Sequence[object]:
         return controller.get_unifi_site_device(site_name=site_name, detailed=detailed, raw=False)
 
     try:
@@ -169,12 +400,18 @@ def fetch_devices(
             )
             return stale_cached
         raise
-    _save_cache(cache_path, devices)
+    if use_cache:
+        _save_cache(cache_path, _serialize_devices_for_cache(devices))
     logger.info("Fetched %d devices", len(devices))
     return devices
 
 
-def fetch_clients(config: Config, *, site: str | None = None) -> Iterable[object]:
+def fetch_clients(
+    config: Config,
+    *,
+    site: str | None = None,
+    use_cache: bool = True,
+) -> Sequence[object]:
     """Fetch active clients from UniFi Controller."""
     try:
         from unifi_controller_api import UnifiAuthenticationError
@@ -183,9 +420,13 @@ def fetch_clients(config: Config, *, site: str | None = None) -> Iterable[object
 
     site_name = site or config.site
     ttl_seconds = _cache_ttl_seconds()
-    cache_path = _cache_dir() / f"clients_{_cache_key(config.url, site_name)}.pkl"
-    cached = _load_cache(cache_path, ttl_seconds)
-    stale_cached, cache_age = _load_cache_with_age(cache_path)
+    cache_path = _cache_dir() / f"clients_{_cache_key(config.url, site_name)}.json"
+    if use_cache and _is_cache_dir_safe(cache_path.parent):
+        cached = _load_cache(cache_path, ttl_seconds)
+        stale_cached, cache_age = _load_cache_with_age(cache_path)
+    else:
+        cached = None
+        stale_cached, cache_age = None, None
     if cached is not None:
         logger.info("Using cached clients (%d)", len(cached))
         return cached
@@ -196,7 +437,7 @@ def fetch_clients(config: Config, *, site: str | None = None) -> Iterable[object
         logger.info("UDM Pro authentication failed, retrying legacy auth")
         controller = _init_controller(config, is_udm_pro=False)
 
-    def _fetch() -> list[object]:
+    def _fetch() -> Sequence[object]:
         return controller.get_unifi_site_client(site_name=site_name, raw=True)
 
     try:
@@ -210,6 +451,7 @@ def fetch_clients(config: Config, *, site: str | None = None) -> Iterable[object
             )
             return stale_cached
         raise
-    _save_cache(cache_path, clients)
+    if use_cache:
+        _save_cache(cache_path, clients)
     logger.info("Fetched %d clients", len(clients))
     return clients