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

"""Sensor platform for AdGuard Control Hub integration."""
import logging
from typing import Any, Optional

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

from . import AdGuardControlHubCoordinator
from .api import AdGuardHomeAPI
from .const import DOMAIN, MANUFACTURER, ICON_STATISTICS, ICON_BLOCKED, ICON_QUERIES, ICON_PERCENTAGE, ICON_CLIENTS

_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 = [
        AdGuardQueriesCounterSensor(coordinator, api),
        AdGuardBlockedCounterSensor(coordinator, api),
        AdGuardBlockingPercentageSensor(coordinator, api),
        AdGuardClientCountSensor(coordinator, api),
        AdGuardProcessingTimeSensor(coordinator, api),
        AdGuardFilteringRulesSensor(coordinator, api),
    ]

    async_add_entities(entities, update_before_add=True)


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

    def __init__(self, coordinator: AdGuardControlHubCoordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator)
        self.api = api
        self._attr_device_info = {
            "identifiers": {(DOMAIN, f"{api.host}:{api.port}")},
            "name": f"AdGuard Control Hub ({api.host})",
            "manufacturer": MANUFACTURER,
            "model": "AdGuard Home",
            "configuration_url": f"{'https' if api.ssl else 'http'}://{api.host}:{api.port}",
        }

    @property
    def available(self) -> bool:
        """Return if sensor is available."""
        return self.coordinator.last_update_success and bool(self.coordinator.statistics)


class AdGuardQueriesCounterSensor(AdGuardBaseSensor):
    """Sensor to track DNS queries count."""

    def __init__(self, coordinator: AdGuardControlHubCoordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_unique_id = f"{api.host}_{api.port}_dns_queries"
        self._attr_name = "AdGuard DNS Queries"
        self._attr_icon = ICON_QUERIES
        self._attr_state_class = SensorStateClass.TOTAL_INCREASING
        self._attr_native_unit_of_measurement = "queries"
        self._attr_entity_category = EntityCategory.DIAGNOSTIC

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

    @property
    def extra_state_attributes(self) -> dict[str, Any]:
        """Return additional state attributes."""
        stats = self.coordinator.statistics
        return {
            "queries_today": stats.get("num_dns_queries_today", 0),
            "replaced_safebrowsing": stats.get("num_replaced_safebrowsing", 0),
            "replaced_parental": stats.get("num_replaced_parental", 0),
            "replaced_safesearch": stats.get("num_replaced_safesearch", 0),
        }


class AdGuardBlockedCounterSensor(AdGuardBaseSensor):
    """Sensor to track blocked queries count."""

    def __init__(self, coordinator: AdGuardControlHubCoordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_unique_id = f"{api.host}_{api.port}_blocked_queries"
        self._attr_name = "AdGuard Blocked Queries"
        self._attr_icon = ICON_BLOCKED
        self._attr_state_class = SensorStateClass.TOTAL_INCREASING
        self._attr_native_unit_of_measurement = "queries"
        self._attr_entity_category = EntityCategory.DIAGNOSTIC

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

    @property
    def extra_state_attributes(self) -> dict[str, Any]:
        """Return additional state attributes."""
        stats = self.coordinator.statistics
        return {
            "blocked_today": stats.get("num_blocked_filtering_today", 0),
            "malware_phishing": stats.get("num_replaced_safebrowsing", 0),
            "adult_websites": stats.get("num_replaced_parental", 0),
        }


class AdGuardBlockingPercentageSensor(AdGuardBaseSensor):
    """Sensor to track blocking percentage."""

    def __init__(self, coordinator: AdGuardControlHubCoordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_unique_id = f"{api.host}_{api.port}_blocking_percentage"
        self._attr_name = "AdGuard Blocking Percentage"
        self._attr_icon = ICON_PERCENTAGE
        self._attr_state_class = SensorStateClass.MEASUREMENT
        self._attr_native_unit_of_measurement = PERCENTAGE
        self._attr_entity_category = EntityCategory.DIAGNOSTIC

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

        if total_queries == 0:
            return 0.0

        percentage = (blocked_queries / total_queries) * 100
        return round(percentage, 2)


class AdGuardClientCountSensor(AdGuardBaseSensor):
    """Sensor to track active clients count."""

    def __init__(self, coordinator: AdGuardControlHubCoordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_unique_id = f"{api.host}_{api.port}_clients_count"
        self._attr_name = "AdGuard Clients Count"
        self._attr_icon = ICON_CLIENTS
        self._attr_state_class = SensorStateClass.MEASUREMENT
        self._attr_native_unit_of_measurement = "clients"
        self._attr_entity_category = EntityCategory.DIAGNOSTIC

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

    @property
    def available(self) -> bool:
        """Return if sensor is available."""
        return self.coordinator.last_update_success

    @property
    def extra_state_attributes(self) -> dict[str, Any]:
        """Return additional state attributes."""
        clients = self.coordinator.clients
        protected_clients = sum(1 for c in clients.values() if c.get("filtering_enabled", True))
        return {
            "protected_clients": protected_clients,
            "unprotected_clients": len(clients) - protected_clients,
            "client_names": list(clients.keys()),
        }


class AdGuardProcessingTimeSensor(AdGuardBaseSensor):
    """Sensor to track average processing time."""

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

    @property
    def native_value(self) -> Optional[float]:
        """Return the state of the sensor."""
        stats = self.coordinator.statistics
        avg_time = stats.get("avg_processing_time", 0)
        return round(avg_time, 2) if avg_time else 0


class AdGuardFilteringRulesSensor(AdGuardBaseSensor):
    """Sensor to track number of filtering rules."""

    def __init__(self, coordinator: AdGuardControlHubCoordinator, api: AdGuardHomeAPI) -> None:
        """Initialize the sensor."""
        super().__init__(coordinator, api)
        self._attr_unique_id = f"{api.host}_{api.port}_filtering_rules"
        self._attr_name = "AdGuard Filtering Rules"
        self._attr_icon = "mdi:filter"
        self._attr_state_class = SensorStateClass.MEASUREMENT
        self._attr_native_unit_of_measurement = "rules"
        self._attr_entity_category = EntityCategory.DIAGNOSTIC

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