unifi-network-maps 1.4.14__py3-none-any.whl → 1.5.0__py3-none-any.whl

unifi_network_maps/cli/render.py

@@ -9,13 +9,11 @@ from ..adapters.config import Config
 from ..adapters.unifi import fetch_clients
 from ..io.export import write_output
 from ..io.mkdocs_assets import write_mkdocs_sidebar_assets
-from ..model.topology import (
-    Device,
-    TopologyResult,
-    build_node_type_map,
-    build_port_map,
-    group_devices_by_type,
-)
+from ..model.classify import classify_device_type
+from ..model.clients import build_node_type_map, collapse_client_edges
+from ..model.edges import build_port_map, group_devices_by_type
+from ..model.topology import Device, Edge, TopologyResult, WanInfo
+from ..model.wan import extract_wan_info
 from ..render.legend import resolve_legend_style
 from ..render.lldp_md import render_lldp_md
 from ..render.mermaid import render_mermaid
@@ -77,6 +75,42 @@ def render_mermaid_output(
     return content
 
 
+def _extract_gateway_wan_info(
+    devices: list[Device],
+    args: argparse.Namespace,
+) -> WanInfo | None:
+    """Extract WAN info from the first gateway device."""
+    for device in devices:
+        if classify_device_type(device) == "gateway":
+            return extract_wan_info(
+                device,
+                wan1_label=getattr(args, "wan_label", None),
+                wan1_isp_speed=getattr(args, "wan_speed", None),
+                wan2_label=getattr(args, "wan2_label", None),
+                wan2_isp_speed=getattr(args, "wan2_speed", None),
+            )
+    return None
+
+
+def _apply_client_clustering(
+    edges: list[Edge],
+    node_types: dict[str, str],
+    layout_mode: str,
+    groups: dict[str, list[str]] | None,
+    group_order: list[str] | None,
+) -> tuple[list[Edge], dict[str, list[str]] | None, list[str] | None]:
+    """Apply client clustering and update groups if needed."""
+    edges, _counts = collapse_client_edges(edges, node_types)
+    if layout_mode == "grouped" and group_order and "client_cluster" not in group_order:
+        group_order = [*group_order, "client_cluster"]
+        if groups is not None:
+            groups = {
+                **groups,
+                "client_cluster": [n for n, t in node_types.items() if t == "client_cluster"],
+            }
+    return edges, groups, group_order
+
+
 def render_svg_output(
     args: argparse.Namespace,
     devices: list[Device],
@@ -89,38 +123,52 @@ def render_svg_output(
 ) -> str:
     edges, _has_tree = select_edges(topology)
     edges, clients = build_edges_with_clients(
-        args,
-        edges,
-        devices,
-        config,
-        site,
-        clients_override=clients_override,
+        args, edges, devices, config, site, clients_override=clients_override
+    )
+    layout_mode = getattr(args, "svg_layout_mode", "physical")
+    options = SvgOptions(width=args.svg_width, height=args.svg_height, layout_mode=layout_mode)
+
+    groups = None
+    group_order = None
+    if layout_mode == "grouped":
+        groups = group_devices_by_type(devices)
+        group_order = ["gateway", "switch", "ap", "other"]
+        if clients:
+            client_names = [c.get("name") or c.get("hostname", "") for c in clients]
+            groups["client"] = [n for n in client_names if n]
+            group_order.append("client")
+
+    node_types = build_node_type_map(
+        devices, clients, client_mode=args.client_scope, only_unifi=args.only_unifi
     )
-    options = SvgOptions(width=args.svg_width, height=args.svg_height)
+
+    if getattr(args, "collapse_clients", False):
+        edges, groups, group_order = _apply_client_clustering(
+            edges, node_types, layout_mode, groups, group_order
+        )
+
+    wan_info = _extract_gateway_wan_info(devices, args)
+
     if args.format == "svg-iso":
-        from ..render.svg import render_svg_isometric
+        from ..render.svg_isometric import render_svg_isometric
 
         return render_svg_isometric(
             edges,
-            node_types=build_node_type_map(
-                devices,
-                clients,
-                client_mode=args.client_scope,
-                only_unifi=args.only_unifi,
-            ),
+            node_types=node_types,
             options=options,
             theme=svg_theme,
+            groups=groups,
+            group_order=group_order,
+            wan_info=wan_info,
         )
     return render_svg(
         edges,
-        node_types=build_node_type_map(
-            devices,
-            clients,
-            client_mode=args.client_scope,
-            only_unifi=args.only_unifi,
-        ),
+        node_types=node_types,
         options=options,
         theme=svg_theme,
+        groups=groups,
+        group_order=group_order,
+        wan_info=wan_info,
     )
 
 
@@ -181,6 +229,7 @@ def render_lldp_format(
     site: str,
     mock_devices: list[object] | None,
     mock_clients: list[object] | None,
+    mock_networks: list[object] | None = None,
 ) -> int:
     try:
         _raw_devices, devices = load_devices_data(
@@ -188,6 +237,7 @@ def render_lldp_format(
             config,
             site,
             raw_devices_override=mock_devices,
+            raw_networks_override=mock_networks,
         )
     except Exception as exc:
         logging.error("Failed to load devices: %s", exc)
@@ -219,6 +269,7 @@ def render_standard_format(
     site: str,
     mock_devices: list[object] | None,
     mock_clients: list[object] | None,
+    mock_networks: list[object] | None = None,
     mermaid_theme: MermaidTheme,
     svg_theme: SvgTheme,
 ) -> int:
@@ -227,6 +278,7 @@ def render_standard_format(
         config=config,
         site=site,
         mock_devices=mock_devices,
+        mock_networks=mock_networks,
     )
     if topology_result is None:
         return 1

unifi_network_maps/cli/runtime.py

@@ -7,19 +7,13 @@ import logging
 from pathlib import Path
 
 from ..adapters.config import Config
-from ..adapters.unifi import fetch_clients, fetch_devices
+from ..adapters.unifi import fetch_clients, fetch_devices, fetch_networks
 from ..io.debug import debug_dump_devices
-from ..model.topology import (
-    ClientPortMap,
-    Device,
-    TopologyResult,
-    build_client_edges,
-    build_client_port_map,
-    build_device_index,
-    build_topology,
-    group_devices_by_type,
-    normalize_devices,
-)
+from ..model.clients import build_client_edges, build_client_port_map
+from ..model.edges import build_topology, enrich_edges_with_active_vlans, group_devices_by_type
+from ..model.topology import ClientPortMap, Device, TopologyResult, build_device_index
+from ..model.topology_coerce import normalize_devices
+from ..model.vlans import build_network_vlan_map
 from ..render.mermaid_theme import MermaidTheme
 from ..render.theme import load_theme
 
@@ -32,6 +26,7 @@ def load_devices_data(
     site: str,
     *,
     raw_devices_override: list[object] | None = None,
+    raw_networks_override: list[object] | None = None,
 ) -> tuple[list[object], list[Device]]:
     if raw_devices_override is None:
         if config is None:
@@ -41,7 +36,19 @@ def load_devices_data(
         )
     else:
         raw_devices = raw_devices_override
-    devices = normalize_devices(raw_devices)
+
+    # Build network-to-VLAN mapping for resolving network IDs in port configs
+    network_vlan_map: dict[str, int] | None = None
+    if raw_networks_override is not None:
+        network_vlan_map = build_network_vlan_map(raw_networks_override)
+    elif config is not None:
+        try:
+            raw_networks = list(fetch_networks(config, site=site, use_cache=not args.no_cache))
+            network_vlan_map = build_network_vlan_map(raw_networks)
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("Failed to fetch networks for VLAN resolution: %s", exc)
+
+    devices = normalize_devices(raw_devices, network_vlan_map)
     if args.debug_dump:
         debug_dump_devices(raw_devices, devices, sample_count=max(0, args.debug_sample))
     return raw_devices, devices
@@ -54,12 +61,14 @@ def build_topology_data(
     *,
     include_ports: bool | None = None,
     raw_devices_override: list[object] | None = None,
+    raw_networks_override: list[object] | None = None,
 ) -> tuple[list[Device], list[str], TopologyResult]:
     _raw_devices, devices = load_devices_data(
         args,
         config,
         site,
         raw_devices_override=raw_devices_override,
+        raw_networks_override=raw_networks_override,
     )
     groups_for_rank = group_devices_by_type(devices)
     gateways = groups_for_rank.get("gateway", [])
@@ -82,6 +91,7 @@ def build_edges_with_clients(
     clients_override: list[object] | None = None,
 ) -> tuple[list, list | None]:
     clients = None
+    client_edges: list = []
     if args.include_clients:
         if clients_override is None:
             if config is None:
@@ -90,14 +100,16 @@ def build_edges_with_clients(
         else:
             clients = clients_override
         device_index = build_device_index(devices)
-        edges = edges + build_client_edges(
+        client_edges = build_client_edges(
             clients,
             device_index,
             include_ports=args.include_ports,
             client_mode=args.client_scope,
             only_unifi=args.only_unifi,
         )
-    return edges, clients
+    # Enrich infrastructure edges with active VLANs from client traffic
+    enriched_edges = enrich_edges_with_active_vlans(edges, client_edges)
+    return enriched_edges + client_edges, clients
 
 
 def select_edges(topology: TopologyResult) -> tuple[list, bool]:
@@ -113,6 +125,7 @@ def load_topology_for_render(
     config: Config | None,
     site: str,
     mock_devices: list[object] | None,
+    mock_networks: list[object] | None = None,
 ) -> tuple[list[Device], TopologyResult] | None:
     try:
         include_ports = True if args.format == "mkdocs" else None
@@ -122,6 +135,7 @@ def load_topology_for_render(
             site,
             include_ports=include_ports,
             raw_devices_override=mock_devices,
+            raw_networks_override=mock_networks,
         )
     except Exception as exc:
         logging.error("Failed to build topology: %s", exc)

unifi_network_maps/io/debug.py

@@ -6,7 +6,8 @@ import json
 import logging
 from collections.abc import Iterable, Sequence
 
-from ..model.topology import Device, group_devices_by_type
+from ..model.edges import group_devices_by_type
+from ..model.topology import Device
 
 logger = logging.getLogger(__name__)
 

unifi_network_maps/io/export.py

@@ -26,16 +26,22 @@ def write_output(
 
 def _write_atomic(path: Path, content: str) -> None:
     path.parent.mkdir(parents=True, exist_ok=True)
-    with tempfile.NamedTemporaryFile(
-        mode="w",
-        encoding="utf-8",
-        dir=str(path.parent),
-        prefix=f".{path.name}.",
-        suffix=".tmp",
-        delete=False,
-    ) as temp_file:
-        temp_file.write(content)
-        temp_file.flush()
-        os.fsync(temp_file.fileno())
-        tmp_path = Path(temp_file.name)
-    tmp_path.replace(path)
+    tmp_path: Path | None = None
+    try:
+        with tempfile.NamedTemporaryFile(
+            mode="w",
+            encoding="utf-8",
+            dir=str(path.parent),
+            prefix=f".{path.name}.",
+            suffix=".tmp",
+            delete=False,
+        ) as temp_file:
+            temp_file.write(content)
+            temp_file.flush()
+            os.fsync(temp_file.fileno())
+            tmp_path = Path(temp_file.name)
+        tmp_path.replace(path)
+    except BaseException:
+        if tmp_path and tmp_path.exists():
+            tmp_path.unlink(missing_ok=True)
+        raise

unifi_network_maps/io/mock_data.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+from pathlib import Path
 
 from .paths import resolve_mock_data_path
 
@@ -15,17 +16,18 @@ def _as_list(value: object, name: str) -> list[object]:
     raise ValueError(f"Mock data field '{name}' must be a list")
 
 
-def load_mock_data(path: str) -> tuple[list[object], list[object]]:
+def load_mock_data(path: str | Path) -> tuple[list[object], list[object], list[object]]:
     resolved = resolve_mock_data_path(path)
     payload = json.loads(resolved.read_text(encoding="utf-8"))
     if not isinstance(payload, dict):
         raise ValueError("Mock data must be a JSON object")
     devices = _as_list(payload.get("devices"), "devices")
     clients = _as_list(payload.get("clients"), "clients")
-    return devices, clients
+    networks = _as_list(payload.get("networks"), "networks")
+    return devices, clients, networks
 
 
-def load_mock_payload(path: str) -> dict[str, list[object] | list[dict[str, object]]]:
+def load_mock_payload(path: str | Path) -> dict[str, list[object] | list[dict[str, object]]]:
     resolved = resolve_mock_data_path(path)
     payload = json.loads(resolved.read_text(encoding="utf-8"))
     if not isinstance(payload, dict):

unifi_network_maps/io/paths.py

@@ -194,8 +194,10 @@ def resolve_output_path(path: str | Path, *, format_name: str | None) -> Path:
 
 
 def resolve_cache_dir(path: str | Path) -> Path:
-    resolved = _resolve_user_path(path)
+    # Check for symlinks on the original path before .resolve() follows them.
+    raw = Path(path).expanduser()
+    _ensure_no_symlink(raw, label="Cache directory")
+    _ensure_no_symlink_in_parents(raw, label="Cache directory")
+    resolved = raw.resolve(strict=False)
     _ensure_within_allowed(resolved, _allowed_roots(), label="Cache directory")
-    _ensure_no_symlink(resolved, label="Cache directory")
-    _ensure_no_symlink_in_parents(resolved, label="Cache directory")
     return resolved

unifi_network_maps/model/classify.py

@@ -0,0 +1,199 @@
+"""Device and client type classification."""
+
+from __future__ import annotations
+
+from .helpers import first_string_field, get_field
+
+# Client device category detection patterns
+_CLIENT_NAME_PATTERNS: dict[str, tuple[str, ...]] = {
+    "camera": ("camera", "cam", "doorbell", "uvc", "protect", "ring", "nest cam", "arlo"),
+    "tv": ("tv", "television", "apple tv", "chromecast", "roku", "fire tv", "shield", "smart tv"),
+    "phone": ("phone", "iphone", "android", "pixel", "galaxy", "mobile", "voip", "handset"),
+    "printer": ("printer", "print", "laserjet", "inkjet", "epson", "canon", "brother", "hp "),
+    "nas": ("nas", "synology", "qnap", "diskstation", "drobo", "freenas", "truenas"),
+    "speaker": ("sonos", "homepod", "echo", "alexa", "google home", "speaker", "soundbar"),
+    "game_console": ("playstation", "ps4", "ps5", "xbox", "nintendo", "switch", "steam deck"),
+    "iot": ("sensor", "thermostat", "nest", "hue", "smart", "zigbee", "z-wave", "iot"),
+}
+
+# OUI/vendor patterns for manufacturer-based detection
+_CLIENT_VENDOR_PATTERNS: dict[str, tuple[str, ...]] = {
+    "camera": ("ubiquiti", "hikvision", "dahua", "axis", "ring", "arlo", "nest", "wyze"),
+    "tv": ("samsung tv", "lg tv", "sony tv", "vizio", "tcl", "roku", "apple tv"),
+    "phone": ("apple", "samsung mobile", "google pixel", "oneplus", "xiaomi"),
+    "printer": ("hp inc", "canon", "epson", "brother", "lexmark", "xerox", "ricoh"),
+    "nas": ("synology", "qnap", "western digital", "seagate", "netgear readynas"),
+    "speaker": ("sonos", "bose", "harman", "bang & olufsen", "denon"),
+    "game_console": ("sony interactive", "microsoft xbox", "nintendo"),
+}
+
+# UniFi product line mappings
+_UNIFI_PRODUCT_CATEGORIES: dict[str, str] = {
+    "protect": "camera",
+    "talk": "phone",
+    "access": "iot",
+    "led": "iot",
+    "connect": "iot",
+}
+
+
+def classify_device_type(device: object) -> str:
+    """Classify a network device into gateway, switch, ap, or other."""
+    raw_type = get_field(device, "type")
+    raw_name = get_field(device, "name")
+    value = raw_type.strip().lower() if isinstance(raw_type, str) else ""
+    if not value:
+        name = raw_name.strip().lower() if isinstance(raw_name, str) else ""
+        if "gateway" in name or name.startswith("gw"):
+            return "gateway"
+        if "switch" in name:
+            return "switch"
+        if "ap" in name:
+            return "ap"
+    if value in {"gateway", "ugw", "usg", "ux", "udm", "udr"}:
+        return "gateway"
+    if value in {"switch", "usw"}:
+        return "switch"
+    if value in {"uap", "ap"} or "ap" in value:
+        return "ap"
+    return "other"
+
+
+def _classify_by_name(name: str) -> str | None:
+    """Classify client by display name heuristics."""
+    name_lower = name.lower()
+    for category, patterns in _CLIENT_NAME_PATTERNS.items():
+        for pattern in patterns:
+            if pattern in name_lower:
+                return category
+    return None
+
+
+def _classify_by_vendor(vendor: str) -> str | None:
+    """Classify client by OUI/vendor name."""
+    vendor_lower = vendor.lower()
+    for category, patterns in _CLIENT_VENDOR_PATTERNS.items():
+        for pattern in patterns:
+            if pattern in vendor_lower:
+                return category
+    return None
+
+
+def _classify_by_unifi_info(ucore: dict[str, object]) -> str | None:
+    """Classify client by UniFi device info (product_line, model)."""
+    product_line = ucore.get("product_line")
+    if isinstance(product_line, str):
+        line_lower = product_line.lower()
+        for line_prefix, category in _UNIFI_PRODUCT_CATEGORIES.items():
+            if line_lower.startswith(line_prefix):
+                return category
+    for key in ("product_shortname", "computed_model", "product_model"):
+        value = ucore.get(key)
+        if isinstance(value, str):
+            value_lower = value.lower()
+            if any(cam in value_lower for cam in ("camera", "doorbell", "uvc", "g4", "g5")):
+                return "camera"
+            if "talk" in value_lower or "phone" in value_lower:
+                return "phone"
+    return None
+
+
+def _client_ucore_info(client: object) -> dict[str, object] | None:
+    """Get UniFi device info from client."""
+    info = get_field(client, "unifi_device_info_from_ucore")
+    if isinstance(info, dict):
+        return info
+    return None
+
+
+def _client_ucore_display_name(client: object) -> str | None:
+    """Get display name from UniFi device info."""
+    ucore = _client_ucore_info(client)
+    if not ucore:
+        return None
+    return first_string_field(ucore, "name", "computed_model", "product_model", "product_shortname")
+
+
+def _client_vendor(client: object) -> str | None:
+    """Get vendor/OUI from client."""
+    return first_string_field(
+        client, "oui", "vendor", "vendor_name", "manufacturer", "manufacturer_name"
+    )
+
+
+def _client_unifi_flag(client: object) -> bool | None:
+    """Check explicit UniFi device flags."""
+    for key in ("is_unifi", "is_unifi_device", "is_ubnt", "is_uap", "is_managed"):
+        value = get_field(client, key)
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, int):
+            return value != 0
+    return None
+
+
+def client_is_unifi(client: object) -> bool:
+    """Determine if a client is a UniFi device."""
+    flag = _client_unifi_flag(client)
+    if flag is not None:
+        return flag
+    ucore = _client_ucore_info(client)
+    if ucore:
+        managed = ucore.get("managed")
+        if isinstance(managed, bool) and managed:
+            return True
+        if isinstance(ucore.get("product_line"), str) and ucore.get("product_line"):
+            return True
+        if isinstance(ucore.get("product_shortname"), str) and ucore.get("product_shortname"):
+            return True
+        for key in ("name", "computed_model", "product_model"):
+            value = ucore.get(key)
+            if isinstance(value, str) and value.strip():
+                return True
+    vendor = _client_vendor(client)
+    if not vendor:
+        return False
+    normalized = vendor.lower()
+    return "ubiquiti" in normalized or "unifi" in normalized
+
+
+def classify_client_type(client: object) -> str:
+    """Classify a client into a device category.
+
+    Detection priority:
+    1. UniFi device info (product_line, model)
+    2. Display name heuristics
+    3. OUI/vendor patterns
+
+    Returns one of: camera, tv, phone, printer, nas, speaker, game_console, iot, client
+    """
+    ucore = _client_ucore_info(client)
+    if ucore:
+        category = _classify_by_unifi_info(ucore)
+        if category:
+            return category
+
+    name = client_display_name(client)
+    if name:
+        category = _classify_by_name(name)
+        if category:
+            return category
+
+    vendor = _client_vendor(client)
+    if vendor:
+        category = _classify_by_vendor(vendor)
+        if category:
+            return category
+
+    return "client"
+
+
+def client_display_name(client: object) -> str | None:
+    """Get display name for a client."""
+    name = first_string_field(client, "name")
+    if name:
+        return name
+    preferred = _client_ucore_display_name(client)
+    if preferred:
+        return preferred
+    return first_string_field(client, "hostname", "mac")