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

"""AdGuard Home API client."""
import asyncio
import logging
from typing import Any, Dict, List, Optional

import aiohttp
from homeassistant.helpers.aiohttp_client import async_get_clientsession

from .const import API_ENDPOINTS

_LOGGER = logging.getLogger(__name__)


class AdGuardHomeError(Exception):
    """Base exception for AdGuard Home errors."""


class AdGuardConnectionError(AdGuardHomeError):
    """Connection error."""


class AdGuardAuthError(AdGuardHomeError):
    """Authentication error."""


class AdGuardTimeoutError(AdGuardHomeError):
    """Timeout error."""


class AdGuardHomeAPI:
    """AdGuard Home API client."""

    def __init__(
        self,
        host: str,
        port: int,
        username: Optional[str] = None,
        password: Optional[str] = None,
        ssl: bool = False,
        verify_ssl: bool = True,
        session: Optional[aiohttp.ClientSession] = None,
        timeout: int = 30,
    ) -> None:
        """Initialize the API client."""
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.ssl = ssl
        self.verify_ssl = verify_ssl
        self.timeout = aiohttp.ClientTimeout(total=timeout)
        self._session = session
        self._auth = None

        if username and password:
            self._auth = aiohttp.BasicAuth(username, password)

    @property
    def base_url(self) -> str:
        """Return the base URL."""
        protocol = "https" if self.ssl else "http"
        return f"{protocol}://{self.host}:{self.port}"

    async def _request(
        self, method: str, endpoint: str, **kwargs
    ) -> Dict[str, Any]:
        """Make a request to the API."""
        url = f"{self.base_url}{endpoint}"

        try:
            async with self._session.request(
                method,
                url,
                auth=self._auth,
                timeout=self.timeout,
                ssl=self.verify_ssl if self.ssl else None,
                **kwargs
            ) as response:
                if response.status == 401:
                    raise AdGuardAuthError("Authentication failed")
                elif response.status == 404:
                    raise AdGuardConnectionError(f"Endpoint not found: {endpoint}")
                elif response.status >= 400:
                    raise AdGuardConnectionError(f"HTTP {response.status}: {response.reason}")

                return await response.json()

        except asyncio.TimeoutError as err:
            raise AdGuardTimeoutError(f"Request timeout for {url}") from err
        except aiohttp.ClientConnectorError as err:
            raise AdGuardConnectionError(f"Connection failed to {url}: {err}") from err
        except aiohttp.ClientError as err:
            raise AdGuardConnectionError(f"Client error for {url}: {err}") from err
        except Exception as err:
            raise AdGuardHomeError(f"Unexpected error for {url}: {err}") from err

    async def test_connection(self) -> bool:
        """Test the connection to AdGuard Home."""
        try:
            await self.get_status()
            return True
        except Exception as err:
            _LOGGER.error("Connection test failed: %s", err)
            return False

    async def get_status(self) -> Dict[str, Any]:
        """Get AdGuard Home status."""
        return await self._request("GET", API_ENDPOINTS["status"])

    async def get_clients(self) -> Dict[str, Any]:
        """Get clients list."""
        return await self._request("GET", API_ENDPOINTS["clients"])

    async def get_statistics(self) -> Dict[str, Any]:
        """Get DNS query statistics."""
        return await self._request("GET", API_ENDPOINTS["stats"])

    async def set_protection(self, enabled: bool) -> None:
        """Enable or disable protection."""
        data = {"enabled": enabled}
        await self._request("POST", API_ENDPOINTS["protection"], json=data)

    async def get_client_by_name(self, name: str) -> Optional[Dict[str, Any]]:
        """Get client by name."""
        clients_data = await self.get_clients()
        for client in clients_data.get("clients", []):
            if client.get("name") == name:
                return client
        return None

    async def update_client_blocked_services(
        self, client_name: str, blocked_services: List[str]
    ) -> None:
        """Update blocked services for a client."""
        client = await self.get_client_by_name(client_name)
        if not client:
            raise AdGuardConnectionError(f"Client '{client_name}' not found")

        # Update client with new blocked services
        client_data = client.copy()
        client_data["blocked_services"] = blocked_services

        await self._request("POST", API_ENDPOINTS["clients_update"], json=client_data)

    async def add_client(self, client_data: Dict[str, Any]) -> None:
        """Add a new client."""
        await self._request("POST", API_ENDPOINTS["clients_add"], json=client_data)

    async def delete_client(self, client_name: str) -> None:
        """Delete a client."""
        data = {"name": client_name}
        await self._request("POST", API_ENDPOINTS["clients_delete"], json=data)