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

"""Config flow for AdGuard Control Hub integration."""
import asyncio
import logging
from typing import Any, Dict, Optional

import voluptuous as vol
from homeassistant import config_entries
from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.data_entry_flow import FlowResult
import homeassistant.helpers.config_validation as cv

from .api import AdGuardHomeAPI, AdGuardConnectionError, AdGuardAuthError
from .const import (
    CONF_SSL,
    CONF_VERIFY_SSL,
    DEFAULT_PORT,
    DEFAULT_SSL,
    DEFAULT_VERIFY_SSL,
    DOMAIN,
)

_LOGGER = logging.getLogger(__name__)

STEP_USER_DATA_SCHEMA = vol.Schema({
    vol.Required(CONF_HOST): cv.string,
    vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
    vol.Optional(CONF_USERNAME): cv.string,
    vol.Optional(CONF_PASSWORD): cv.string,
    vol.Optional(CONF_SSL, default=DEFAULT_SSL): cv.boolean,
    vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): cv.boolean,
})


async def validate_input(hass, data: Dict[str, Any]) -> Dict[str, Any]:
    """Validate the user input allows us to connect."""
    # Normalize host
    host = data[CONF_HOST].strip()
    if not host:
        raise InvalidHost("Host cannot be empty")

    # Remove protocol if provided
    if host.startswith(("http://", "https://")):
        host = host.split("://", 1)[1]
    data[CONF_HOST] = host

    # Validate port
    port = data[CONF_PORT]
    if not (1 <= port <= 65535):
        raise InvalidPort("Port must be between 1 and 65535")

    # Create session with appropriate SSL settings
    session = async_get_clientsession(hass, data.get(CONF_VERIFY_SSL, True))

    # Create API instance
    api = AdGuardHomeAPI(
        host=host,
        port=port,
        username=data.get(CONF_USERNAME),
        password=data.get(CONF_PASSWORD),
        ssl=data.get(CONF_SSL, False),
        session=session,
        timeout=10,  # 10 second timeout for setup
    )

    # Test the connection
    try:
        if not await api.test_connection():
            raise CannotConnect("Failed to connect to AdGuard Home")

        # Get additional server info if possible
        try:
            status = await api.get_status()
            version = status.get("version", "unknown")
            dns_port = status.get("dns_port", "N/A")

            return {
                "title": f"AdGuard Control Hub ({host})",
                "version": version,
                "dns_port": dns_port,
                "host": host,
            }
        except Exception as err:
            _LOGGER.warning("Could not get server status, but connection works: %s", err)
            return {
                "title": f"AdGuard Control Hub ({host})",
                "version": "unknown",
                "dns_port": "N/A",
                "host": host,
            }

    except AdGuardAuthError as err:
        _LOGGER.error("Authentication failed: %s", err)
        raise InvalidAuth from err
    except AdGuardConnectionError as err:
        _LOGGER.error("Connection failed: %s", err)
        if "timeout" in str(err).lower():
            raise Timeout from err
        raise CannotConnect from err
    except asyncio.TimeoutError as err:
        _LOGGER.error("Connection timeout: %s", err)
        raise Timeout from err
    except Exception as err:
        _LOGGER.exception("Unexpected error during validation: %s", err)
        raise CannotConnect from err


class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
    """Handle a config flow for AdGuard Control Hub."""

    VERSION = 1
    CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_POLL

    async def async_step_user(
        self, user_input: Optional[Dict[str, Any]] = None
    ) -> FlowResult:
        """Handle the initial step."""
        errors: Dict[str, str] = {}

        if user_input is not None:
            try:
                info = await validate_input(self.hass, user_input)

                # Create unique ID based on host and port
                unique_id = f"{info['host']}:{user_input[CONF_PORT]}"
                await self.async_set_unique_id(unique_id)
                self._abort_if_unique_id_configured()

                return self.async_create_entry(
                    title=info["title"],
                    data=user_input,
                )

            except CannotConnect:
                errors["base"] = "cannot_connect"
            except InvalidAuth:
                errors["base"] = "invalid_auth"
            except InvalidHost:
                errors[CONF_HOST] = "invalid_host"
            except InvalidPort:
                errors[CONF_PORT] = "invalid_port"
            except Timeout:
                errors["base"] = "timeout"
            except Exception:
                _LOGGER.exception("Unexpected exception during config flow")
                errors["base"] = "unknown"

        return self.async_show_form(
            step_id="user",
            data_schema=STEP_USER_DATA_SCHEMA,
            errors=errors,
        )

    async def async_step_import(self, import_info: Dict[str, Any]) -> FlowResult:
        """Handle configuration import."""
        return await self.async_step_user(import_info)

    @staticmethod
    def async_get_options_flow(config_entry):
        """Get the options flow for this handler."""
        return OptionsFlowHandler(config_entry)


class OptionsFlowHandler(config_entries.OptionsFlow):
    """Handle options flow for AdGuard Control Hub."""

    def __init__(self, config_entry: config_entries.ConfigEntry) -> None:
        """Initialize options flow."""
        self.config_entry = config_entry

    async def async_step_init(
        self, user_input: Optional[Dict[str, Any]] = None
    ) -> FlowResult:
        """Handle options flow."""
        if user_input is not None:
            return self.async_create_entry(title="", data=user_input)

        options_schema = vol.Schema({
            vol.Optional(
                "scan_interval",
                default=self.config_entry.options.get("scan_interval", 30),
            ): vol.All(vol.Coerce(int), vol.Range(min=10, max=300)),
            vol.Optional(
                "timeout",
                default=self.config_entry.options.get("timeout", 10),
            ): vol.All(vol.Coerce(int), vol.Range(min=5, max=60)),
        })

        return self.async_show_form(
            step_id="init",
            data_schema=options_schema,
        )


# Custom exceptions
class CannotConnect(Exception):
    """Error to indicate we cannot connect."""


class InvalidAuth(Exception):
    """Error to indicate there is invalid auth."""


class InvalidHost(Exception):
    """Error to indicate invalid host."""


class InvalidPort(Exception):
    """Error to indicate invalid port."""


class Timeout(Exception):
    """Error to indicate connection timeout."""