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

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

from homeassistant.components.sensor import SensorEntity, SensorStateClass
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import PERCENTAGE
from homeassistant.core import HomeAssistant
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

_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),
    ]

    async_add_entities(entities)


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

    def __init__(self, coordinator: AdGuardControlHubCoordinator, api: AdGuardHomeAPI):
        """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",
        }


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

    def __init__(self, coordinator: AdGuardControlHubCoordinator, api: AdGuardHomeAPI):
        """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_STATISTICS
        self._attr_state_class = SensorStateClass.TOTAL_INCREASING
        self._attr_native_unit_of_measurement = "queries"

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


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

    def __init__(self, coordinator: AdGuardControlHubCoordinator, api: AdGuardHomeAPI):
        """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_STATISTICS
        self._attr_state_class = SensorStateClass.TOTAL_INCREASING
        self._attr_native_unit_of_measurement = "queries"

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


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

    def __init__(self, coordinator: AdGuardControlHubCoordinator, api: AdGuardHomeAPI):
        """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_STATISTICS
        self._attr_state_class = SensorStateClass.MEASUREMENT
        self._attr_native_unit_of_measurement = PERCENTAGE

    @property
    def native_value(self):
        """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):
        """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_STATISTICS
        self._attr_state_class = SensorStateClass.MEASUREMENT
        self._attr_native_unit_of_measurement = "clients"

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