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

unifi_network_maps/io/mock_generate.py

@@ -0,0 +1,299 @@
+"""Generate mock UniFi data using Faker."""
+
+from __future__ import annotations
+
+import json
+import random
+from dataclasses import dataclass, field
+from typing import Any
+
+from faker import Faker
+
+
+@dataclass(frozen=True)
+class MockOptions:
+    seed: int = 1337
+    switch_count: int = 1
+    ap_count: int = 2
+    wired_client_count: int = 2
+    wireless_client_count: int = 2
+
+
+@dataclass
+class _MockState:
+    fake: Faker
+    rng: random.Random
+    used_macs: set[str] = field(default_factory=set)
+    used_ips: set[str] = field(default_factory=set)
+    used_names: set[str] = field(default_factory=set)
+    used_rooms: set[str] = field(default_factory=set)
+    core_port_next: int = 2
+
+
+def generate_mock_payload(options: MockOptions) -> dict[str, list[dict[str, Any]]]:
+    state = _build_state(options.seed)
+    devices, core_switch, aps = _build_devices(options, state)
+    clients = _build_clients(options, state, core_switch, aps)
+    return {"devices": devices, "clients": clients}
+
+
+def mock_payload_json(options: MockOptions) -> str:
+    payload = generate_mock_payload(options)
+    return json.dumps(payload, indent=2, sort_keys=True)
+
+
+def _build_state(seed: int) -> _MockState:
+    fake = Faker("en_US")
+    fake.seed_instance(seed)
+    rng = random.Random(seed)
+    return _MockState(fake=fake, rng=rng)
+
+
+def _build_devices(
+    options: MockOptions, state: _MockState
+) -> tuple[list[dict[str, Any]], dict[str, Any], list[dict[str, Any]]]:
+    gateway = _build_gateway(state)
+    core_switch = _build_core_switch(state)
+    _link_gateway_to_switch(state, gateway, core_switch)
+    access_switches = _build_access_switches(options.switch_count - 1, state, core_switch)
+    aps = _build_aps(options.ap_count, state, core_switch)
+    devices = [gateway, core_switch] + access_switches + aps
+    return devices, core_switch, aps
+
+
+def _build_gateway(state: _MockState) -> dict[str, Any]:
+    device = _device_base(state, "Cloud Gateway", "udm", "UniFi Dream Machine Pro", "UDM-Pro")
+    _add_port(device, 9, poe_enabled=False, rng=state.rng)
+    return device
+
+
+def _build_core_switch(state: _MockState) -> dict[str, Any]:
+    return _device_base(state, "Core Switch", "usw", "UniFi Switch 24 PoE", "USW-24-PoE")
+
+
+def _build_access_switches(
+    count: int, state: _MockState, core_switch: dict[str, Any]
+) -> list[dict[str, Any]]:
+    switches = []
+    for _ in range(max(0, count)):
+        room = _unique_room(state)
+        name = _unique_name(state, f"Switch {room}")
+        device = _device_base(state, name, "usw", "UniFi Switch Lite 8 PoE", "USW-Lite-8-PoE")
+        _link_core_device(state, core_switch, device, poe_enabled=False)
+        switches.append(device)
+    return switches
+
+
+def _build_aps(count: int, state: _MockState, core_switch: dict[str, Any]) -> list[dict[str, Any]]:
+    aps = []
+    for _ in range(max(0, count)):
+        room = _unique_room(state)
+        name = _unique_name(state, f"AP {room}")
+        device = _device_base(state, name, "uap", "UniFi AP 6 Lite", "U6-Lite")
+        _add_port(device, 1, poe_enabled=True, rng=state.rng)
+        _link_core_device(state, core_switch, device, poe_enabled=True)
+        aps.append(device)
+    return aps
+
+
+def _build_clients(
+    options: MockOptions,
+    state: _MockState,
+    core_switch: dict[str, Any],
+    aps: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+    clients = []
+    clients.extend(_build_wired_clients(options.wired_client_count, state, core_switch))
+    clients.extend(_build_wireless_clients(options.wireless_client_count, state, aps))
+    return clients
+
+
+def _build_wired_clients(
+    count: int, state: _MockState, core_switch: dict[str, Any]
+) -> list[dict[str, Any]]:
+    clients = []
+    for _ in range(max(0, count)):
+        port_idx = _next_core_port(state)
+        _add_port(core_switch, port_idx, poe_enabled=False, rng=state.rng)
+        name = _unique_client_name(state)
+        clients.append(
+            {
+                "name": name,
+                "is_wired": True,
+                "sw_mac": core_switch["mac"],
+                "sw_port": port_idx,
+            }
+        )
+    return clients
+
+
+def _build_wireless_clients(
+    count: int, state: _MockState, aps: list[dict[str, Any]]
+) -> list[dict[str, Any]]:
+    if not aps:
+        return []
+    clients = []
+    for idx in range(max(0, count)):
+        ap = aps[idx % len(aps)]
+        name = _unique_client_name(state)
+        clients.append(
+            {
+                "name": name,
+                "is_wired": False,
+                "ap_mac": ap["mac"],
+                "ap_port": 1,
+            }
+        )
+    return clients
+
+
+def _device_base(
+    state: _MockState, name: str, dev_type: str, model_name: str, model: str
+) -> dict[str, Any]:
+    version = _pick_version(state, dev_type)
+    return {
+        "name": name,
+        "model_name": model_name,
+        "model": model,
+        "mac": _unique_mac(state),
+        "ip": _unique_ip(state),
+        "type": dev_type,
+        "version": version,
+        "port_table": [],
+        "lldp_info": [],
+    }
+
+
+def _pick_version(state: _MockState, dev_type: str) -> str:
+    versions = {
+        "udm": ["3.1.0", "3.1.1"],
+        "usw": ["7.0.0", "7.1.2"],
+        "uap": ["6.6.55", "6.7.10"],
+    }
+    return state.rng.choice(versions.get(dev_type, ["1.0.0"]))
+
+
+def _link_gateway_to_switch(
+    state: _MockState, gateway: dict[str, Any], core_switch: dict[str, Any]
+) -> None:
+    _add_port(core_switch, 1, poe_enabled=False, rng=state.rng)
+    _add_lldp_link(
+        gateway,
+        core_switch,
+        local_port=9,
+        remote_port=1,
+        remote_name=core_switch["name"],
+    )
+    _add_lldp_link(
+        core_switch,
+        gateway,
+        local_port=1,
+        remote_port=9,
+        remote_name=gateway["name"],
+    )
+
+
+def _link_core_device(
+    state: _MockState,
+    core_switch: dict[str, Any],
+    device: dict[str, Any],
+    *,
+    poe_enabled: bool,
+) -> None:
+    port_idx = _next_core_port(state)
+    _add_port(core_switch, port_idx, poe_enabled=poe_enabled, rng=state.rng)
+    _add_lldp_link(
+        core_switch,
+        device,
+        local_port=port_idx,
+        remote_port=1,
+        remote_name=device["name"],
+    )
+    _add_lldp_link(
+        device,
+        core_switch,
+        local_port=1,
+        remote_port=port_idx,
+        remote_name=core_switch["name"],
+    )
+
+
+def _add_lldp_link(
+    source: dict[str, Any],
+    dest: dict[str, Any],
+    *,
+    local_port: int,
+    remote_port: int,
+    remote_name: str,
+) -> None:
+    source["lldp_info"].append(
+        {
+            "chassis_id": dest["mac"],
+            "port_id": f"Port {remote_port}",
+            "port_desc": remote_name,
+            "local_port_name": f"Port {local_port}",
+            "local_port_idx": local_port,
+        }
+    )
+
+
+def _add_port(
+    device: dict[str, Any], port_idx: int, *, poe_enabled: bool, rng: random.Random
+) -> None:
+    entry = {
+        "port_idx": port_idx,
+        "name": f"Port {port_idx}",
+        "ifname": f"eth{max(0, port_idx - 1)}",
+        "poe_enable": poe_enabled,
+        "port_poe": poe_enabled,
+    }
+    if poe_enabled:
+        entry["poe_power"] = round(rng.uniform(4.5, 7.5), 1)
+    device["port_table"].append(entry)
+
+
+def _unique_mac(state: _MockState) -> str:
+    return _unique_value(state, state.fake.mac_address, state.used_macs)
+
+
+def _unique_ip(state: _MockState) -> str:
+    return _unique_value(state, state.fake.ipv4_private, state.used_ips)
+
+
+def _unique_room(state: _MockState) -> str:
+    def _room() -> str:
+        return state.fake.word().title()
+
+    return _unique_value(state, _room, state.used_rooms)
+
+
+def _unique_name(state: _MockState, candidate: str) -> str:
+    if candidate not in state.used_names:
+        state.used_names.add(candidate)
+        return candidate
+    suffix = state.rng.randint(2, 99)
+    value = f"{candidate} {suffix}"
+    state.used_names.add(value)
+    return value
+
+
+def _unique_client_name(state: _MockState) -> str:
+    def _name() -> str:
+        return state.fake.hostname()
+
+    return _unique_value(state, _name, state.used_names)
+
+
+def _unique_value(state: _MockState, generator, used: set[str]) -> str:
+    for _ in range(100):
+        value = str(generator())
+        if value not in used:
+            used.add(value)
+            return value
+    raise ValueError("Failed to generate a unique value")
+
+
+def _next_core_port(state: _MockState) -> int:
+    port_idx = state.core_port_next
+    state.core_port_next += 1
+    return port_idx

unifi_network_maps/model/lldp.py

@@ -17,18 +17,32 @@ class LLDPEntry:
 
 
 def coerce_lldp(entry: object) -> LLDPEntry:
-    chassis_id = getattr(entry, "chassis_id", None) or getattr(entry, "chassisId", None)
-    port_id = getattr(entry, "port_id", None) or getattr(entry, "portId", None)
-    port_desc = (
-        getattr(entry, "port_desc", None)
-        or getattr(entry, "portDesc", None)
-        or getattr(entry, "port_descr", None)
-        or getattr(entry, "portDescr", None)
-    )
-    local_port_name = getattr(entry, "local_port_name", None) or getattr(
-        entry, "localPortName", None
-    )
-    local_port_idx = getattr(entry, "local_port_idx", None) or getattr(entry, "localPortIdx", None)
+    if isinstance(entry, dict):
+        chassis_id = entry.get("chassis_id") or entry.get("chassisId")
+        port_id = entry.get("port_id") or entry.get("portId")
+        port_desc = (
+            entry.get("port_desc")
+            or entry.get("portDesc")
+            or entry.get("port_descr")
+            or entry.get("portDescr")
+        )
+        local_port_name = entry.get("local_port_name") or entry.get("localPortName")
+        local_port_idx = entry.get("local_port_idx") or entry.get("localPortIdx")
+    else:
+        chassis_id = getattr(entry, "chassis_id", None) or getattr(entry, "chassisId", None)
+        port_id = getattr(entry, "port_id", None) or getattr(entry, "portId", None)
+        port_desc = (
+            getattr(entry, "port_desc", None)
+            or getattr(entry, "portDesc", None)
+            or getattr(entry, "port_descr", None)
+            or getattr(entry, "portDescr", None)
+        )
+        local_port_name = getattr(entry, "local_port_name", None) or getattr(
+            entry, "localPortName", None
+        )
+        local_port_idx = getattr(entry, "local_port_idx", None) or getattr(
+            entry, "localPortIdx", None
+        )
     if not chassis_id or not port_id:
         raise ValueError("LLDP entry missing chassis_id or port_id")
     return LLDPEntry(

unifi_network_maps/model/topology.py

@@ -6,7 +6,7 @@ import logging
 from collections import deque
 from collections.abc import Iterable
 from dataclasses import dataclass, field
-from typing import Protocol, TypeAlias
+from typing import Protocol
 
 from .labels import compose_port_label, order_edge_names
 from .lldp import LLDPEntry, coerce_lldp, local_port_label
@@ -19,6 +19,7 @@ logger = logging.getLogger(__name__)
 class Device:
     name: str
     model_name: str
+    model: str
     mac: str
     ip: str
     type: str
@@ -74,17 +75,22 @@ class PortInfo:
     port_idx: int | None
     name: str | None
     ifname: str | None
+    speed: int | None
+    aggregation_group: str | None
     port_poe: bool
     poe_enable: bool
     poe_good: bool
     poe_power: float | None
 
 
-PortMap: TypeAlias = dict[tuple[str, str], str]
-PoeMap: TypeAlias = dict[tuple[str, str], bool]
+type PortMap = dict[tuple[str, str], str]
+type PoeMap = dict[tuple[str, str], bool]
+type ClientPortMap = dict[str, list[tuple[int, str]]]
 
 
 def _get_attr(obj: object, name: str) -> object | None:
+    if isinstance(obj, dict):
+        return obj.get(name)
     return getattr(obj, name, None)
 
 
@@ -123,6 +129,44 @@ def _as_int(value: object | None) -> int | None:
     return None
 
 
+def _as_group_id(value: object | None) -> str | None:
+    if value is None:
+        return None
+    if isinstance(value, bool):
+        return None
+    if isinstance(value, int):
+        return str(value)
+    if isinstance(value, str):
+        return value.strip() or None
+    return None
+
+
+def _aggregation_group(port_entry: object) -> object | None:
+    keys = (
+        "aggregation_group",
+        "aggregation_id",
+        "aggregate_id",
+        "agg_id",
+        "lag_id",
+        "lag_group",
+        "link_aggregation_group",
+        "link_aggregation_id",
+        "aggregate",
+        "aggregated_by",
+    )
+    if isinstance(port_entry, dict):
+        for key in keys:
+            value = port_entry.get(key)
+            if value not in (None, "", False):
+                return value
+        return None
+    for key in keys:
+        value = _get_attr(port_entry, key)
+        if value not in (None, "", False):
+            return value
+    return None
+
+
 def _resolve_port_idx_from_lldp(lldp_entry: LLDPEntry, port_table: list[PortInfo]) -> int | None:
     if lldp_entry.local_port_idx is not None:
         return lldp_entry.local_port_idx
@@ -153,6 +197,8 @@ def _port_info_from_entry(port_entry: object) -> PortInfo:
         port_idx = port_entry.get("port_idx") or port_entry.get("portIdx")
         name = port_entry.get("name")
         ifname = port_entry.get("ifname")
+        speed = port_entry.get("speed")
+        aggregation_group = _aggregation_group(port_entry)
         port_poe = _as_bool(port_entry.get("port_poe"))
         poe_enable = _as_bool(port_entry.get("poe_enable"))
         poe_good = _as_bool(port_entry.get("poe_good"))
@@ -161,6 +207,8 @@ def _port_info_from_entry(port_entry: object) -> PortInfo:
         port_idx = _get_attr(port_entry, "port_idx") or _get_attr(port_entry, "portIdx")
         name = _get_attr(port_entry, "name")
         ifname = _get_attr(port_entry, "ifname")
+        speed = _get_attr(port_entry, "speed")
+        aggregation_group = _aggregation_group(port_entry)
         port_poe = _as_bool(_get_attr(port_entry, "port_poe"))
         poe_enable = _as_bool(_get_attr(port_entry, "poe_enable"))
         poe_good = _as_bool(_get_attr(port_entry, "poe_good"))
@@ -169,6 +217,8 @@ def _port_info_from_entry(port_entry: object) -> PortInfo:
         port_idx=_as_int(port_idx),
         name=str(name) if isinstance(name, str) and name.strip() else None,
         ifname=str(ifname) if isinstance(ifname, str) and ifname.strip() else None,
+        speed=_as_int(speed),
+        aggregation_group=_as_group_id(aggregation_group),
         port_poe=port_poe,
         poe_enable=poe_enable,
         poe_good=poe_good,
@@ -243,6 +293,7 @@ def _uplink_info(device: DeviceLike) -> tuple[UplinkInfo | None, UplinkInfo | No
 def coerce_device(device: DeviceLike) -> Device:
     name = _get_attr(device, "name")
     model_name = _get_attr(device, "model_name") or _get_attr(device, "model")
+    model = _get_attr(device, "model")
     mac = _get_attr(device, "mac")
     ip = _get_attr(device, "ip") or _get_attr(device, "ip_address")
     dev_type = _get_attr(device, "type") or _get_attr(device, "device_type")
@@ -268,6 +319,7 @@ def coerce_device(device: DeviceLike) -> Device:
     return Device(
         name=str(name),
         model_name=str(model_name or ""),
+        model=str(model or ""),
         mac=str(mac),
         ip=str(ip or ""),
         type=str(dev_type or ""),
@@ -528,6 +580,62 @@ def build_edges(
     return edges
 
 
+def build_port_map(devices: Iterable[Device], *, only_unifi: bool = True) -> PortMap:
+    ordered_devices = sorted(devices, key=lambda item: (item.name.lower(), item.mac.lower()))
+    index = build_device_index(ordered_devices)
+    device_by_name = {device.name: device for device in ordered_devices}
+    raw_links: list[tuple[str, str]] = []
+    seen: set[frozenset[str]] = set()
+    port_map: PortMap = {}
+    poe_map: PoeMap = {}
+
+    devices_with_lldp_edges = _collect_lldp_links(
+        ordered_devices,
+        index,
+        port_map,
+        poe_map,
+        raw_links,
+        seen,
+        only_unifi=only_unifi,
+    )
+    _collect_uplink_links(
+        ordered_devices,
+        devices_with_lldp_edges,
+        index,
+        device_by_name,
+        port_map,
+        poe_map,
+        raw_links,
+        seen,
+        include_ports=True,
+        only_unifi=only_unifi,
+    )
+    return port_map
+
+
+def build_client_port_map(
+    devices: Iterable[Device],
+    clients: Iterable[object],
+    *,
+    client_mode: str,
+) -> ClientPortMap:
+    device_index = build_device_index(devices)
+    port_map: ClientPortMap = {}
+    for client in clients:
+        if not _client_matches_mode(client, client_mode):
+            continue
+        name = _client_display_name(client)
+        uplink_mac = _client_uplink_mac(client)
+        uplink_port = _client_uplink_port(client)
+        if not name or not uplink_mac or uplink_port is None:
+            continue
+        device_name = device_index.get(_normalize_mac(uplink_mac))
+        if not device_name:
+            continue
+        port_map.setdefault(device_name, []).append((uplink_port, name))
+    return port_map
+
+
 def _collect_lldp_links(
     devices: list[Device],
     index: dict[str, str],