diff --git a/custom_components/technitiumdns/__init__.py b/custom_components/technitiumdns/__init__.py index 3d74773..29ca3eb 100644 --- a/custom_components/technitiumdns/__init__.py +++ b/custom_components/technitiumdns/__init__.py @@ -11,7 +11,6 @@ _LOGGER = logging.getLogger(__name__) - async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up TechnitiumDNS from a config entry.""" hass.data.setdefault(DOMAIN, {}) @@ -26,8 +25,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: device_registry.async_get_or_create( config_entry_id=entry.entry_id, identifiers={(DOMAIN, entry.entry_id)}, - name=entry.data["server_name"], manufacturer="Technitium", + name=entry.data["server_name"], model="DNS Server", ) @@ -37,7 +36,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: ) return True - async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload a config entry.""" await hass.config_entries.async_unload_platforms( diff --git a/custom_components/technitiumdns/button.py b/custom_components/technitiumdns/button.py index 7b5ef37..dbf869a 100644 --- a/custom_components/technitiumdns/button.py +++ b/custom_components/technitiumdns/button.py @@ -1,15 +1,13 @@ import logging from homeassistant.components.button import ButtonEntity -from homeassistant.core import HomeAssistant -from homeassistant.helpers.update_coordinator import CoordinatorEntity +from homeassistant.helpers.entity import DeviceInfo from .const import DOMAIN, AD_BLOCKING_DURATION_OPTIONS from .api import TechnitiumDNSApi _LOGGER = logging.getLogger(__name__) - -async def async_setup_entry(hass: HomeAssistant, entry, async_add_entities): +async def async_setup_entry(hass, entry, async_add_entities): """Set up TechnitiumDNS button entities based on a config entry.""" config_entry = hass.data[DOMAIN][entry.entry_id] api = config_entry["api"] @@ -21,11 +19,7 @@ async def async_setup_entry(hass: HomeAssistant, entry, async_add_entities): # Define the buttons using the sorted durations buttons = [ TechnitiumDNSButton( - api, - AD_BLOCKING_DURATION_OPTIONS[duration], - duration, - server_name, - entry.entry_id, + api, AD_BLOCKING_DURATION_OPTIONS[duration], duration, server_name, entry.entry_id ) for duration in sorted_durations ] @@ -33,18 +27,10 @@ async def async_setup_entry(hass: HomeAssistant, entry, async_add_entities): # Add entities async_add_entities(buttons) - class TechnitiumDNSButton(ButtonEntity): """Representation of a TechnitiumDNS button.""" - def __init__( - self, - api: TechnitiumDNSApi, - name: str, - duration: int, - server_name: str, - entry_id: str, - ): + def __init__(self, api: TechnitiumDNSApi, name: str, duration: int, server_name: str, entry_id: str): """Initialize the button.""" self._api = api self._attr_name = f"{name} ({server_name})" @@ -55,19 +41,16 @@ async def async_press(self) -> None: """Handle the button press.""" try: await self._api.temporary_disable_blocking(self._duration) - _LOGGER.info( - f"Ad blocking disabled for {self._duration} minutes on {self._attr_name}" - ) + _LOGGER.info(f"Ad blocking disabled for {self._duration} minutes on {self._attr_name}") except Exception as e: _LOGGER.error(f"Failed to disable ad blocking: {e}") @property def device_info(self): """Return device information for this entity.""" - return { - "identifiers": {(DOMAIN, self._entry_id)}, - "name": self._attr_name, - "manufacturer": "Technitium", - "model": "DNS Server", - "entry_type": "service", - } + return DeviceInfo( + identifiers={(DOMAIN, self._entry_id)}, + name=self._attr_name, + manufacturer="Technitium", + model="DNS Server", + ) diff --git a/custom_components/technitiumdns/sensor.py b/custom_components/technitiumdns/sensor.py index 6edb80b..3f8d06d 100644 --- a/custom_components/technitiumdns/sensor.py +++ b/custom_components/technitiumdns/sensor.py @@ -7,7 +7,7 @@ UpdateFailed, CoordinatorEntity, ) -from homeassistant.helpers import device_registry as dr +from homeassistant.helpers.entity import DeviceInfo from .const import DOMAIN, SENSOR_TYPES from .api import TechnitiumDNSApi @@ -16,7 +16,6 @@ SCAN_INTERVAL = timedelta(minutes=1) - async def async_setup_entry(hass, entry, async_add_entities): """Set up the TechnitiumDNS sensor based on a config entry.""" config_entry = hass.data[DOMAIN][entry.entry_id] @@ -27,18 +26,12 @@ async def async_setup_entry(hass, entry, async_add_entities): coordinator = TechnitiumDNSCoordinator(hass, api, stats_duration) await coordinator.async_config_entry_first_refresh() - device_registry = dr.async_get(hass) - device = device_registry.async_get_device(identifiers={(DOMAIN, entry.entry_id)}) - sensors = [] for sensor_type in SENSOR_TYPES: - sensors.append( - TechnitiumDNSSensor(coordinator, sensor_type, server_name, device.id) - ) + sensors.append(TechnitiumDNSSensor(coordinator, sensor_type, server_name, entry.entry_id)) async_add_entities(sensors, True) - class TechnitiumDNSCoordinator(DataUpdateCoordinator): """Class to manage fetching TechnitiumDNS data.""" @@ -52,52 +45,25 @@ async def _async_update_data(self): """Update data via library.""" try: _LOGGER.debug("Fetching data from TechnitiumDNS API") - Technitiumdns_statistics = await self.api.get_statistics( - self.stats_duration - ) - Technitiumdns_top_clients = await self.api.get_top_clients( - self.stats_duration - ) - Technitiumdns_top_domains = await self.api.get_top_domains( - self.stats_duration - ) - Technitiumdns_top_blocked_domains = await self.api.get_top_blocked_domains( - self.stats_duration - ) + Technitiumdns_statistics = await self.api.get_statistics(self.stats_duration) + Technitiumdns_top_clients = await self.api.get_top_clients(self.stats_duration) + Technitiumdns_top_domains = await self.api.get_top_domains(self.stats_duration) + Technitiumdns_top_blocked_domains = await self.api.get_top_blocked_domains(self.stats_duration) Technitiumdns_update_info = await self.api.check_update() # Add more logging to debug empty response issue - _LOGGER.debug( - "Technitiumdns_statistics response content: %s", - Technitiumdns_statistics, - ) - _LOGGER.debug( - "Technitiumdns_top_clients response content: %s", - Technitiumdns_top_clients, - ) - _LOGGER.debug( - "Technitiumdns_top_domains response content: %s", - Technitiumdns_top_domains, - ) - _LOGGER.debug( - "Technitiumdns_top_blocked_domains response content: %s", - Technitiumdns_top_blocked_domains, - ) - _LOGGER.debug( - "Technitiumdns_update_info response content: %s", - Technitiumdns_update_info, - ) - - Technitiumdns_stats = Technitiumdns_statistics.get("response", {}).get( - "stats", {} - ) + _LOGGER.debug("Technitiumdns_statistics response content: %s", Technitiumdns_statistics) + _LOGGER.debug("Technitiumdns_top_clients response content: %s", Technitiumdns_top_clients) + _LOGGER.debug("Technitiumdns_top_domains response content: %s", Technitiumdns_top_domains) + _LOGGER.debug("Technitiumdns_top_blocked_domains response content: %s", Technitiumdns_top_blocked_domains) + _LOGGER.debug("Technitiumdns_update_info response content: %s", Technitiumdns_update_info) + + Technitiumdns_stats = Technitiumdns_statistics.get("response", {}).get("stats", {}) data = { "queries": Technitiumdns_stats.get("totalQueries"), "blocked_queries": Technitiumdns_stats.get("totalBlocked"), "clients": Technitiumdns_stats.get("totalClients"), - "update_available": Technitiumdns_update_info.get("response", {}).get( - "updateAvailable" - ), + "update_available": Technitiumdns_update_info.get("response", {}).get("updateAvailable"), "no_error": Technitiumdns_stats.get("totalNoError"), "server_failure": Technitiumdns_stats.get("totalServerFailure"), "nx_domain": Technitiumdns_stats.get("totalNxDomain"), @@ -113,28 +79,13 @@ async def _async_update_data(self): "allow_list_zones": Technitiumdns_stats.get("allowListZones"), "block_list_zones": Technitiumdns_stats.get("blockListZones"), "top_clients": "\n".join( - [ - f"{client['name']} ({client['hits']})" - for client in Technitiumdns_top_clients.get("response", {}).get( - "topClients", [] - )[:5] - ] + [f"{client['name']} ({client['hits']})" for client in Technitiumdns_top_clients.get("response", {}).get("topClients", [])[:5]] ), "top_domains": "\n".join( - [ - f"{domain['name']} ({domain['hits']})" - for domain in Technitiumdns_top_domains.get("response", {}).get( - "topDomains", [] - )[:5] - ] + [f"{domain['name']} ({domain['hits']})" for domain in Technitiumdns_top_domains.get("response", {}).get("topDomains", [])[:5]] ), "top_blocked_domains": "\n".join( - [ - f"{domain['name']} ({domain['hits']})" - for domain in Technitiumdns_top_blocked_domains.get( - "response", {} - ).get("topBlockedDomains", [])[:5] - ] + [f"{domain['name']} ({domain['hits']})" for domain in Technitiumdns_top_blocked_domains.get("response", {}).get("topBlockedDomains", [])[:5]] ), } _LOGGER.debug("Data combined: %s", data) @@ -143,19 +94,16 @@ async def _async_update_data(self): _LOGGER.error("Error fetching data: %s", err) raise UpdateFailed(f"Error fetching data: {err}") - class TechnitiumDNSSensor(CoordinatorEntity, SensorEntity): """Representation of a TechnitiumDNS sensor.""" - def __init__(self, coordinator, sensor_type, server_name, device_id): + def __init__(self, coordinator, sensor_type, server_name, entry_id): """Initialize the sensor.""" super().__init__(coordinator) self._sensor_type = sensor_type self._server_name = server_name - self._device_id = device_id - self._name = ( - f"Technitiumdns_{SENSOR_TYPES[sensor_type]['name']} ({server_name})" - ) + self._entry_id = entry_id + self._name = f"Technitiumdns_{SENSOR_TYPES[sensor_type]['name']} ({server_name})" @property def name(self): @@ -170,9 +118,7 @@ def state(self): # Ensure the state value is within the allowable length if isinstance(state_value, str) and len(state_value) > 255: - _LOGGER.error( - "State value for %s exceeds 255 characters", self._sensor_type - ) + _LOGGER.error("State value for %s exceeds 255 characters", self._sensor_type) return state_value[:255] if isinstance(state_value, (list, dict)): @@ -198,10 +144,10 @@ def should_poll(self): @property def device_info(self): - """Return the device info.""" - return { - "identifiers": {(DOMAIN, self._device_id)}, - "name": self._server_name, - "manufacturer": "Technitium", - "model": "DNS Server", - } + """Return device information for this entity.""" + return DeviceInfo( + identifiers={(DOMAIN, self._entry_id)}, + name=self._server_name, + manufacturer="Technitium", + model="DNS Server", + ) diff --git a/custom_components/technitiumdns/switch.py b/custom_components/technitiumdns/switch.py index 26b9e8e..4e4c37f 100644 --- a/custom_components/technitiumdns/switch.py +++ b/custom_components/technitiumdns/switch.py @@ -1,35 +1,28 @@ import logging from homeassistant.components.switch import SwitchEntity -from homeassistant.core import HomeAssistant -from homeassistant.helpers.update_coordinator import CoordinatorEntity +from homeassistant.helpers.entity import DeviceInfo from .const import DOMAIN, AD_BLOCKING_SWITCH from .api import TechnitiumDNSApi _LOGGER = logging.getLogger(__name__) - -async def async_setup_entry(hass: HomeAssistant, entry, async_add_entities): +async def async_setup_entry(hass, entry, async_add_entities): """Set up TechnitiumDNS switch entities based on a config entry.""" config_entry = hass.data[DOMAIN][entry.entry_id] api = config_entry["api"] server_name = config_entry["server_name"] # Define the switch - switches = [ - TechnitiumDNSSwitch(api, AD_BLOCKING_SWITCH, server_name, entry.entry_id) - ] + switches = [TechnitiumDNSSwitch(api, AD_BLOCKING_SWITCH, server_name, entry.entry_id)] # Add entities async_add_entities(switches) - class TechnitiumDNSSwitch(SwitchEntity): """Representation of a TechnitiumDNS switch.""" - def __init__( - self, api: TechnitiumDNSApi, name: str, server_name: str, entry_id: str - ): + def __init__(self, api: TechnitiumDNSApi, name: str, server_name: str, entry_id: str): """Initialize the switch.""" self._api = api self._attr_name = f"{name} ({server_name})" @@ -55,9 +48,7 @@ async def _fetch_state(self): try: response = await self._api.get_dns_settings() self._is_on = response["response"].get("enableBlocking", False) - _LOGGER.info( - f"Fetched ad blocking state: {self._is_on} for {self._attr_name}" - ) + _LOGGER.info(f"Fetched ad blocking state: {self._is_on} for {self._attr_name}") self.async_write_ha_state() except Exception as e: _LOGGER.error(f"Failed to fetch ad blocking state: {e}") @@ -85,10 +76,9 @@ async def async_turn_off(self, **kwargs): @property def device_info(self): """Return device information for this entity.""" - return { - "identifiers": {(DOMAIN, self._entry_id)}, - "name": self._attr_name, - "manufacturer": "Technitium", - "model": "DNS Server", - "entry_type": "service", - } + return DeviceInfo( + identifiers={(DOMAIN, self._entry_id)}, + name=self._attr_name, + manufacturer="Technitium", + model="DNS Server", + )