adguard-control-hub/custom_components/adguard_hub/sensor.py

"""AdGuard Control Hub sensor platform."""
import logging
from typing import Any, Dict, List, Optional

from homeassistant.components.sensor import (
    SensorEntity,
    SensorDeviceClass,
    SensorStateClass,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from homeassistant.helpers.entity import DeviceInfo, EntityCategory
from homeassistant.const import PERCENTAGE, UnitOfTime

from .api import AdGuardHomeAPI
from .const import DOMAIN, MANUFACTURER

_LOGGER = logging.getLogger(__name__)


async def async_setup_entry(
    hass: HomeAssistant,
    config_entry: ConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up AdGuard Control Hub sensor platform."""
    coordinator = hass.data[DOMAIN][config_entry.entry_id]["coordinator"]
    api = hass.data[DOMAIN][config_entry.entry_id]["api"]

    entities: List[SensorEntity] = []

    # Add main sensors
    entities.extend([
        AdGuardQueriesCounterSensor(coordinator, api),
        AdGuardBlockedCounterSensor(coordinator, api),
        AdGuardBlockingPercentageSensor(coordinator, api),
        AdGuardClientsCountSensor(coordinator, api),
        AdGuardProcessingTimeSensor(coordinator, api),
        AdGuardFilteringRulesSensor(coordinator, api),
        AdGuardUpstreamServersSensor(coordinator, api),
        AdGuardVersionSensor(coordinator, api),
    ])

    async_add_entities(entities)


class AdGuardBaseSensor(CoordinatorEntity, SensorEntity):
    """Base AdGuard sensor."""

    def __init__(self, coordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator)
        self.api = api

    @property
    def device_info(self) -> DeviceInfo:
        """Return device info."""
        return DeviceInfo(
            identifiers={(DOMAIN, "adguard_home")},
            name="AdGuard Home",
            manufacturer=MANUFACTURER,
            model="AdGuard Home",
            configuration_url=self.api.base_url,
        )


class AdGuardQueriesCounterSensor(AdGuardBaseSensor):
    """AdGuard DNS queries counter sensor."""

    def __init__(self, coordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_name = "AdGuard DNS Queries"
        self._attr_unique_id = f"{DOMAIN}_dns_queries"
        self._attr_device_class = SensorDeviceClass.ENUM
        self._attr_state_class = SensorStateClass.TOTAL_INCREASING
        self._attr_icon = "mdi:dns"

    @property
    def native_value(self) -> Optional[int]:
        """Return the state of the sensor."""
        return self.coordinator.statistics.get("num_dns_queries", 0)


class AdGuardBlockedCounterSensor(AdGuardBaseSensor):
    """AdGuard blocked queries counter sensor."""

    def __init__(self, coordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_name = "AdGuard Blocked Queries"
        self._attr_unique_id = f"{DOMAIN}_blocked_queries"
        self._attr_device_class = SensorDeviceClass.ENUM
        self._attr_state_class = SensorStateClass.TOTAL_INCREASING
        self._attr_icon = "mdi:shield-check"

    @property
    def native_value(self) -> Optional[int]:
        """Return the state of the sensor."""
        return self.coordinator.statistics.get("num_blocked_filtering", 0)


class AdGuardBlockingPercentageSensor(AdGuardBaseSensor):
    """AdGuard blocking percentage sensor."""

    def __init__(self, coordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_name = "AdGuard Blocking Percentage"
        self._attr_unique_id = f"{DOMAIN}_blocking_percentage"
        self._attr_device_class = SensorDeviceClass.ENUM
        self._attr_state_class = SensorStateClass.MEASUREMENT
        self._attr_native_unit_of_measurement = PERCENTAGE
        self._attr_icon = "mdi:percent"

    @property
    def native_value(self) -> Optional[float]:
        """Return the state of the sensor."""
        total_queries = self.coordinator.statistics.get("num_dns_queries", 0)
        blocked_queries = self.coordinator.statistics.get("num_blocked_filtering", 0)

        if total_queries > 0:
            return round((blocked_queries / total_queries) * 100, 2)
        return 0.0


class AdGuardClientsCountSensor(AdGuardBaseSensor):
    """AdGuard clients count sensor."""

    def __init__(self, coordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_name = "AdGuard Clients Count"
        self._attr_unique_id = f"{DOMAIN}_clients_count"
        self._attr_device_class = SensorDeviceClass.ENUM
        self._attr_state_class = SensorStateClass.MEASUREMENT
        self._attr_icon = "mdi:account-multiple"
        self._attr_entity_category = EntityCategory.DIAGNOSTIC

    @property
    def native_value(self) -> int:
        """Return the state of the sensor."""
        return len(self.coordinator.clients)


class AdGuardProcessingTimeSensor(AdGuardBaseSensor):
    """AdGuard average processing time sensor."""

    def __init__(self, coordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_name = "AdGuard Average Processing Time"
        self._attr_unique_id = f"{DOMAIN}_avg_processing_time"
        self._attr_device_class = SensorDeviceClass.DURATION
        self._attr_state_class = SensorStateClass.MEASUREMENT
        self._attr_native_unit_of_measurement = UnitOfTime.MILLISECONDS
        self._attr_icon = "mdi:speedometer"
        self._attr_entity_category = EntityCategory.DIAGNOSTIC

    @property
    def native_value(self) -> Optional[float]:
        """Return the state of the sensor."""
        return self.coordinator.statistics.get("avg_processing_time", 0.0)


class AdGuardFilteringRulesSensor(AdGuardBaseSensor):
    """AdGuard filtering rules count sensor."""

    def __init__(self, coordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_name = "AdGuard Filtering Rules"
        self._attr_unique_id = f"{DOMAIN}_filtering_rules"
        self._attr_device_class = SensorDeviceClass.ENUM
        self._attr_state_class = SensorStateClass.MEASUREMENT
        self._attr_icon = "mdi:filter"
        self._attr_entity_category = EntityCategory.DIAGNOSTIC

    @property
    def native_value(self) -> Optional[int]:
        """Return the state of the sensor."""
        return self.coordinator.protection_status.get("num_filtering_rules", 0)


class AdGuardUpstreamServersSensor(AdGuardBaseSensor):
    """AdGuard upstream servers sensor."""

    def __init__(self, coordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_name = "AdGuard Upstream Servers"
        self._attr_unique_id = f"{DOMAIN}_upstream_servers"
        self._attr_icon = "mdi:server-network"
        self._attr_entity_category = EntityCategory.DIAGNOSTIC

    @property
    def native_value(self) -> str:
        """Return the state of the sensor."""
        servers = self.coordinator.protection_status.get("dns_addresses", [])
        return ", ".join(servers) if servers else "Unknown"


class AdGuardVersionSensor(AdGuardBaseSensor):
    """AdGuard version sensor."""

    def __init__(self, coordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_name = "AdGuard Version"
        self._attr_unique_id = f"{DOMAIN}_version"
        self._attr_icon = "mdi:information"
        self._attr_entity_category = EntityCategory.DIAGNOSTIC

    @property
    def native_value(self) -> str:
        """Return the state of the sensor."""
        return self.coordinator.protection_status.get("version", "Unknown")