repairs.py

Go to the documentation of this file.

"""unifiprotect.repairs."""
 
from __future__ import annotations
 
from typing import cast
 
from uiprotect import ProtectApiClient
from uiprotect.data import Bootstrap, Camera, ModelType
from uiprotect.data.types import FirmwareReleaseChannel
import voluptuous as vol
 
from homeassistant import data_entry_flow
from homeassistant.components.repairs import ConfirmRepairFlow, RepairsFlow
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import issue_registry as ir
 
from .const import CONF_ALLOW_EA
from .data import UFPConfigEntry, async_get_data_for_entry_id
from .utils import async_create_api_client
 
 
class ProtectRepair(RepairsFlow):
    """Handler for an issue fixing flow."""
 
    _api: ProtectApiClient
    _entry: UFPConfigEntry
 
    def __init__(self, *, api: ProtectApiClient, entry: UFPConfigEntry) -> None:
        """Create flow."""
 
        self._api_api = api
        self._entry_entry = entry
        super().__init__()
 
    @callback
    def _async_get_placeholders(self) -> dict[str, str]:
        issue_registry = ir.async_get(self.hass)
        description_placeholders = {}
        if issue := issue_registry.async_get_issue(self.handler, self.issue_id):
            description_placeholders = issue.translation_placeholders or {}
            if issue.learn_more_url:
                description_placeholders["learn_more"] = issue.learn_more_url
 
        return description_placeholders
 
 
class EAConfirmRepair(ProtectRepair):
    """Handler for an issue fixing flow."""
 
    async def async_step_init(
        self, user_input: dict[str, str] | None = None
    ) -> data_entry_flow.FlowResult:
        """Handle the first step of a fix flow."""
 
        return await self.async_step_startasync_step_start()
 
    async def async_step_start(
        self, user_input: dict[str, str] | None = None
    ) -> data_entry_flow.FlowResult:
        """Handle the confirm step of a fix flow."""
        if user_input is None:
            placeholders = self._async_get_placeholders_async_get_placeholders()
            return self.async_show_formasync_show_form(
                step_id="start",
                data_schema=vol.Schema({}),
                description_placeholders=placeholders,
            )
 
        nvr = await self._api_api.get_nvr()
        if nvr.release_channel != FirmwareReleaseChannel.RELEASE:
            return await self.async_step_confirmasync_step_confirm()
        await self.hass.config_entries.async_reload(self._entry_entry.entry_id)
        return self.async_create_entryasync_create_entry(data={})
 
    async def async_step_confirm(
        self, user_input: dict[str, str] | None = None
    ) -> data_entry_flow.FlowResult:
        """Handle the confirm step of a fix flow."""
        if user_input is not None:
            options = dict(self._entry_entry.options)
            options[CONF_ALLOW_EA] = True
            self.hass.config_entries.async_update_entry(self._entry_entry, options=options)
            return self.async_create_entryasync_create_entry(data={})
 
        placeholders = self._async_get_placeholders_async_get_placeholders()
        return self.async_show_formasync_show_form(
            step_id="confirm",
            data_schema=vol.Schema({}),
            description_placeholders=placeholders,
        )
 
 
class CloudAccountRepair(ProtectRepair):
    """Handler for an issue fixing flow."""
 
    async def async_step_init(
        self, user_input: dict[str, str] | None = None
    ) -> data_entry_flow.FlowResult:
        """Handle the first step of a fix flow."""
 
        return await self.async_step_confirmasync_step_confirm()
 
    async def async_step_confirm(
        self, user_input: dict[str, str] | None = None
    ) -> data_entry_flow.FlowResult:
        """Handle the first step of a fix flow."""
 
        if user_input is None:
            placeholders = self._async_get_placeholders_async_get_placeholders()
            return self.async_show_formasync_show_form(
                step_id="confirm",
                data_schema=vol.Schema({}),
                description_placeholders=placeholders,
            )
 
        self._entry_entry.async_start_reauth(self.hass)
        return self.async_create_entryasync_create_entry(data={})
 
 
class RTSPRepair(ProtectRepair):
    """Handler for an issue fixing flow."""
 
    _camera_id: str
    _camera: Camera | None
    _bootstrap: Bootstrap | None
 
    def __init__(
        self,
        *,
        api: ProtectApiClient,
        entry: UFPConfigEntry,
        camera_id: str,
    ) -> None:
        """Create flow."""
 
        super().__init__(api=api, entry=entry)
        self._camera_id_camera_id = camera_id
        self._bootstrap_bootstrap = None
        self._camera_camera = None
 
    @callback
    def _async_get_placeholders(self) -> dict[str, str]:
        description_placeholders = super()._async_get_placeholders()
        if self._camera_camera is not None:
            description_placeholders["camera"] = self._camera_camera.display_name
 
        return description_placeholders
 
    async def _get_boostrap(self) -> Bootstrap:
        if self._bootstrap_bootstrap is None:
            self._bootstrap_bootstrap = await self._api_api.get_bootstrap()
 
        return self._bootstrap_bootstrap
 
    async def _get_camera(self) -> Camera:
        if self._camera_camera is None:
            bootstrap = await self._get_boostrap_get_boostrap()
            self._camera_camera = bootstrap.cameras.get(self._camera_id_camera_id)
            assert self._camera_camera is not None
        return self._camera_camera
 
    async def _enable_rtsp(self) -> None:
        camera = await self._get_camera_get_camera()
        bootstrap = await self._get_boostrap_get_boostrap()
        user = bootstrap.users.get(bootstrap.auth_user_id)
        if not user or not camera.can_write(user):
            return
 
        channel = camera.channels[0]
        channel.is_rtsp_enabled = True
        await self._api_api.update_device(
            ModelType.CAMERA, camera.id, {"channels": camera.unifi_dict()["channels"]}
        )
 
    async def async_step_init(
        self, user_input: dict[str, str] | None = None
    ) -> data_entry_flow.FlowResult:
        """Handle the first step of a fix flow."""
 
        return await self.async_step_startasync_step_start()
 
    async def async_step_start(
        self, user_input: dict[str, str] | None = None
    ) -> data_entry_flow.FlowResult:
        """Handle the first step of a fix flow."""
 
        if user_input is None:
            # make sure camera object is loaded for placeholders
            await self._get_camera_get_camera()
            placeholders = self._async_get_placeholders_async_get_placeholders_async_get_placeholders()
            return self.async_show_formasync_show_form(
                step_id="start",
                data_schema=vol.Schema({}),
                description_placeholders=placeholders,
            )
 
        updated_camera = await self._api_api.get_camera(self._camera_id_camera_id)
        if not any(c.is_rtsp_enabled for c in updated_camera.channels):
            await self._enable_rtsp_enable_rtsp()
 
        updated_camera = await self._api_api.get_camera(self._camera_id_camera_id)
        if any(c.is_rtsp_enabled for c in updated_camera.channels):
            await self.hass.config_entries.async_reload(self._entry_entry.entry_id)
            return self.async_create_entryasync_create_entry(data={})
        return await self.async_step_confirmasync_step_confirm()
 
    async def async_step_confirm(
        self, user_input: dict[str, str] | None = None
    ) -> data_entry_flow.FlowResult:
        """Handle the confirm step of a fix flow."""
        if user_input is not None:
            return self.async_create_entryasync_create_entry(data={})
 
        placeholders = self._async_get_placeholders_async_get_placeholders_async_get_placeholders()
        return self.async_show_formasync_show_form(
            step_id="confirm",
            data_schema=vol.Schema({}),
            description_placeholders=placeholders,
        )
 
 
@callback
def _async_get_or_create_api_client(
    hass: HomeAssistant, entry: ConfigEntry
) -> ProtectApiClient:
    """Get or create an API client."""
    if data := async_get_data_for_entry_id(hass, entry.entry_id):
        return data.api
    return async_create_api_client(hass, entry)
 
 
async def async_create_fix_flow(
    hass: HomeAssistant,
    issue_id: str,
    data: dict[str, str | int | float | None] | None,
) -> RepairsFlow:
    """Create flow."""
    if (
        data is not None
        and "entry_id" in data
        and (entry := hass.config_entries.async_get_entry(cast(str, data["entry_id"])))
    ):
        api = _async_get_or_create_api_client(hass, entry)
        if issue_id == "ea_channel_warning":
            return EAConfirmRepair(api=api, entry=entry)
        if issue_id == "cloud_user":
            return CloudAccountRepair(api=api, entry=entry)
        if issue_id.startswith("rtsp_disabled_"):
            return RTSPRepair(
                api=api, entry=entry, camera_id=cast(str, data["camera_id"])
            )
    return ConfirmRepairFlow()