"""Traffic module.""" from __future__ import annotations from enum import StrEnum import logging from typing import Any, Final from asusrouter.const import ( AR_CALL_GET_STATE, AR_CALL_TRANSLATE_STATE, HTTPStatus, ) from asusrouter.modules.endpoint import ( EndpointTools, EndpointType, get_request_type, ) from asusrouter.modules.source import ARDataSource from asusrouter.registry import ARCallableRegistry as ARCallReg from asusrouter.tools.converters import flatten_dict, safe_bool_nn from asusrouter.tools.enum import FromStrMixin from asusrouter.tools.identifiers import MacAddress from asusrouter.tools.readers import read_units_data_rate from asusrouter.tools.types import ARCallableType, ARCallbackType from asusrouter.tools.units import UnitOfDataRate from asusrouter.tools.writers import dict_to_request _LOGGER = logging.getLogger(__name__) # Initialize data rate readers once _read_kibps = read_units_data_rate(UnitOfDataRate.KIBIBIT_PER_SECOND) _read_Mibps = read_units_data_rate(UnitOfDataRate.MEBIBIT_PER_SECOND) # noqa: N816 _TRANSLATION_TABLE: dict[str, tuple[str, ARCallableType]] = { # Data rate / speed (current) "data_rx": ("rx_speed", _read_kibps), "data_tx": ("tx_speed", _read_kibps), # Data rate / speed (average) "data_avg_rx": ("rx_speed_avg", _read_kibps), "data_avg_tx": ("tx_speed_avg", _read_kibps), # Physical layer rate / speed "phy_rx": ("phy_rx", _read_Mibps), "phy_tx": ("phy_tx", _read_Mibps), # Status code "error_status": ("status", HTTPStatus.from_value), } class ARTrafficType(FromStrMixin, StrEnum): """Traffic type.""" UNKNOWN = "unknown" BACKHAUL = "backhaul" ETHERNET = "ethernet" WIFI = "wifi" AR_TRAFFIC_SOURCE: Final[dict[ARTrafficType, EndpointType]] = { ARTrafficType.BACKHAUL: EndpointTools.TRAFFIC_BACKHAUL, ARTrafficType.ETHERNET: EndpointTools.TRAFFIC_ETHERNET, ARTrafficType.WIFI: EndpointTools.TRAFFIC_WIFI, } class ARTrafficSource(ARDataSource): """Traffic data source.""" def __init__(self, target: Any, type: Any = None) -> None: """Initialize the traffic source.""" super().__init__() self._target: MacAddress | None = None self._type: ARTrafficType = ARTrafficType.UNKNOWN # Try to set the target and type from input self.target = target self.type = type def __repr__(self) -> str: """Representation of the traffic source.""" return f"" @property def target(self) -> MacAddress | None: """Get the target.""" return self._target @target.setter def target(self, value: Any) -> None: """Set the target.""" try: self._target = MacAddress(value) except ValueError: self._target = None @property def type(self) -> ARTrafficType: """Get the traffic type.""" return self._type @type.setter def type(self, type: ARTrafficType) -> None: """Set the traffic type.""" self._type = ARTrafficType.from_value(type) class ARTrafficSourceEthernet(ARTrafficSource): """Ethernet traffic source. This source provides data for: - LAN if `bh_flag` is `False` - Wired backhaul to the parent (if `bh_flag` is `True`) """ def __init__( self, target: Any, type: Any = None, bh_flag: Any = None ) -> None: """Initialize the Ethernet traffic source.""" super().__init__(target, type) self.type = ARTrafficType.ETHERNET self._bh_flag: bool self.bh_flag = bh_flag @property def bh_flag(self) -> bool: """Get the backhaul flag.""" return self._bh_flag @bh_flag.setter def bh_flag(self, value: Any) -> None: """Set the backhaul flag.""" self._bh_flag = safe_bool_nn(value) class ARTrafficSourceBetween(ARTrafficSource): """Traffic source between two points.""" def __init__( self, target: Any, type: Any = None, towards: Any = None ) -> None: """Initialize the traffic source between two points.""" super().__init__(target, type) self._towards: MacAddress | None self.towards = towards @property def towards(self) -> MacAddress | None: """Get the target MAC address.""" return self._towards @towards.setter def towards(self, value: Any) -> None: """Set the target MAC address.""" try: self._towards = MacAddress(value) except ValueError: self._towards = None class ARTrafficSourceWiFi(ARTrafficSourceBetween): """WiFi traffic source. Traffic of the router / node (`target`) on the specific wireless band (`towards`) defined by its MAC address. """ def __init__(self, target: Any, towards: Any) -> None: """Initialize the WiFi traffic source.""" super().__init__(target, type=ARTrafficType.WIFI, towards=towards) self.type = ARTrafficType.WIFI class ARTrafficSourceBackhaul(ARTrafficSourceBetween): """Backhaul traffic source. Traffic from the router / node (`target`) to the specific node (`towards`) defined by its MAC address. """ def __init__(self, target: Any, towards: Any) -> None: """Initialize the backhaul traffic source.""" super().__init__(target, type=ARTrafficType.BACKHAUL, towards=towards) self.type = ARTrafficType.BACKHAUL def _check_state(source: Any) -> None: """Check the state of the traffic source. This helper allows checking that all the required arguments are present and valid. """ # Is a traffic source if not isinstance(source, ARTrafficSource): raise TypeError(f"Expected `ARTrafficSource`, got `{type(source)}`") # Has target and type set if source.target is None: raise ValueError("Traffic source must have a `target` property set") if source.type == ARTrafficType.UNKNOWN: raise ValueError("Traffic source must have a `type` property set") # Has a `toward` property if required if isinstance(source, ARTrafficSourceBetween) and source.towards is None: raise ValueError( f"Traffic source of type `{source.type}` requires " "a `towards` property set" ) async def get_state( callback: ARCallbackType, source: ARTrafficSource, **kwargs: Any, ) -> Any: """Get the state of the traffic source.""" # Check the input source _check_state(source) # Get the correct traffic endpoint endpoint: EndpointType | None = AR_TRAFFIC_SOURCE.get(source.type, None) if endpoint is None: raise ValueError( f"Cannot find endpoint for traffic type `{source.type}`" ) _LOGGER.debug("Getting state for traffic source: %s", source) arguments: dict[str, Any] = { "node_mac": source.target, } # Add backhaul flag if required if isinstance(source, ARTrafficSourceEthernet): arguments["is_bh"] = source.bh_flag # Add towards channel if required if isinstance(source, ARTrafficSourceBackhaul): arguments["sta_mac"] = source.towards if isinstance(source, ARTrafficSourceWiFi): arguments["band_mac"] = source.towards # Call the callback to get the state request = dict_to_request( arguments, request_type=get_request_type(endpoint) ) return await callback(endpoint=endpoint, request=request) def translate_state( data: dict[str, Any], **kwargs: Any, ) -> dict[str, Any]: """Translate the traffic data to a simple format.""" result: dict[str, Any] = {} simplified_data = flatten_dict(data) if not simplified_data: return {} for key, value in simplified_data.items(): if key in _TRANSLATION_TABLE: new_key, transform = _TRANSLATION_TABLE[key] result[new_key] = transform(value) continue result[key] = value return result # Register callables for traffic sources calls: dict[str, ARCallableType] = { AR_CALL_GET_STATE: get_state, AR_CALL_TRANSLATE_STATE: translate_state, } ARCallReg.register(ARTrafficSourceEthernet, **calls) ARCallReg.register(ARTrafficSourceWiFi, **calls) ARCallReg.register(ARTrafficSourceBackhaul, **calls)