core/homeassistant/components/go2rtc/__init__.py

"""The go2rtc component."""

from __future__ import annotations

from dataclasses import dataclass
import logging
from secrets import token_hex
import shutil

from aiohttp import BasicAuth, ClientSession, UnixConnector
from aiohttp.client_exceptions import ClientConnectionError, ServerConnectionError
from awesomeversion import AwesomeVersion
from go2rtc_client import Go2RtcRestClient
from go2rtc_client.exceptions import Go2RtcClientError, Go2RtcVersionError
from go2rtc_client.ws import (
    Go2RtcWsClient,
    ReceiveMessages,
    WebRTCAnswer,
    WebRTCCandidate,
    WebRTCOffer,
    WsError,
)
import voluptuous as vol
from webrtc_models import RTCIceCandidateInit

from homeassistant.components.camera import (
    Camera,
    CameraWebRTCProvider,
    WebRTCAnswer as HAWebRTCAnswer,
    WebRTCCandidate as HAWebRTCCandidate,
    WebRTCError,
    WebRTCMessage,
    WebRTCSendMessage,
    async_register_webrtc_provider,
    get_dynamic_camera_stream_settings,
)
from homeassistant.components.default_config import DOMAIN as DEFAULT_CONFIG_DOMAIN
from homeassistant.components.stream import Orientation
from homeassistant.config_entries import SOURCE_SYSTEM, ConfigEntry
from homeassistant.const import (
    CONF_PASSWORD,
    CONF_URL,
    CONF_USERNAME,
    EVENT_HOMEASSISTANT_STOP,
)
from homeassistant.core import Event, HomeAssistant, callback
from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError
from homeassistant.helpers import (
    config_validation as cv,
    discovery_flow,
    issue_registry as ir,
)
from homeassistant.helpers.aiohttp_client import (
    async_create_clientsession,
    async_get_clientsession,
)
from homeassistant.helpers.typing import ConfigType
from homeassistant.util.hass_dict import HassKey
from homeassistant.util.package import is_docker_env

from .const import (
    CONF_DEBUG_UI,
    DEBUG_UI_URL_MESSAGE,
    DOMAIN,
    HA_MANAGED_UNIX_SOCKET,
    HA_MANAGED_URL,
    RECOMMENDED_VERSION,
)
from .server import Server

_LOGGER = logging.getLogger(__name__)

_FFMPEG = "ffmpeg"
_AUTH = "auth"


def _validate_auth(config: dict) -> dict:
    """Validate that username and password are only set when a URL is configured or when debug UI is enabled."""
    auth_exists = CONF_USERNAME in config
    debug_ui_enabled = config.get(CONF_DEBUG_UI, False)

    if debug_ui_enabled and not auth_exists:
        raise vol.Invalid("Username and password must be set when debug_ui is true")

    if auth_exists and CONF_URL not in config and not debug_ui_enabled:
        raise vol.Invalid(
            "Username and password can only be set when a URL is configured or debug_ui is true"
        )

    return config


CONFIG_SCHEMA = vol.Schema(
    {
        DOMAIN: vol.All(
            vol.Schema(
                {
                    vol.Exclusive(CONF_URL, DOMAIN, DEBUG_UI_URL_MESSAGE): cv.url,
                    vol.Exclusive(
                        CONF_DEBUG_UI, DOMAIN, DEBUG_UI_URL_MESSAGE
                    ): cv.boolean,
                    vol.Inclusive(CONF_USERNAME, _AUTH): vol.All(
                        cv.string, vol.Length(min=1)
                    ),
                    vol.Inclusive(CONF_PASSWORD, _AUTH): vol.All(
                        cv.string, vol.Length(min=1)
                    ),
                }
            ),
            _validate_auth,
        )
    },
    extra=vol.ALLOW_EXTRA,
)

_DATA_GO2RTC: HassKey[Go2RtcConfig] = HassKey(DOMAIN)
_RETRYABLE_ERRORS = (ClientConnectionError, ServerConnectionError)
type Go2RtcConfigEntry = ConfigEntry[WebRTCProvider]


async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
    """Set up WebRTC."""
    url: str | None = None
    username: str | None = None
    password: str | None = None

    if DOMAIN not in config and DEFAULT_CONFIG_DOMAIN not in config:
        await _remove_go2rtc_entries(hass)
        return True

    domain_config = config.get(DOMAIN, {})
    username = domain_config.get(CONF_USERNAME)
    password = domain_config.get(CONF_PASSWORD)

    if not (configured_by_user := DOMAIN in config) or not (
        url := domain_config.get(CONF_URL)
    ):
        if not is_docker_env():
            if not configured_by_user:
                # Remove config entry if it exists
                await _remove_go2rtc_entries(hass)
                return True
            _LOGGER.warning("Go2rtc URL required in non-docker installs")
            return False
        if not (binary := await _get_binary(hass)):
            _LOGGER.error("Could not find go2rtc docker binary")
            return False

        # Generate random credentials when not provided to secure the server
        if not username or not password:
            username = token_hex()
            password = token_hex()
            _LOGGER.debug("Generated random credentials for go2rtc server")

        auth = BasicAuth(username, password)
        # HA will manage the binary
        # Manually created session (not using the helper) needs to be closed manually
        # See on_stop listener below
        session = ClientSession(
            connector=UnixConnector(path=HA_MANAGED_UNIX_SOCKET), auth=auth
        )
        server = Server(
            hass,
            binary,
            session,
            enable_ui=domain_config.get(CONF_DEBUG_UI, False),
            username=username,
            password=password,
        )
        try:
            await server.start()
        except Exception:  # noqa: BLE001
            _LOGGER.warning("Could not start go2rtc server", exc_info=True)
            return False

        async def on_stop(event: Event) -> None:
            await server.stop()
            await session.close()

        hass.bus.async_listen(EVENT_HOMEASSISTANT_STOP, on_stop)

        url = HA_MANAGED_URL
    elif username and password:
        # Create session with BasicAuth if credentials are provided
        auth = BasicAuth(username, password)
        session = async_create_clientsession(hass, auth=auth)
    else:
        session = async_get_clientsession(hass)

    hass.data[_DATA_GO2RTC] = Go2RtcConfig(url, session)
    discovery_flow.async_create_flow(
        hass, DOMAIN, context={"source": SOURCE_SYSTEM}, data={}
    )
    return True


async def _remove_go2rtc_entries(hass: HomeAssistant) -> None:
    """Remove go2rtc config entries, if any."""
    for entry in hass.config_entries.async_entries(DOMAIN):
        await hass.config_entries.async_remove(entry.entry_id)


async def async_setup_entry(hass: HomeAssistant, entry: Go2RtcConfigEntry) -> bool:
    """Set up go2rtc from a config entry."""

    config = hass.data[_DATA_GO2RTC]
    url = config.url
    session = config.session
    client = Go2RtcRestClient(session, url)
    # Validate the server URL
    try:
        version = await client.validate_server_version()
        if version < AwesomeVersion(RECOMMENDED_VERSION):
            ir.async_create_issue(
                hass,
                DOMAIN,
                "recommended_version",
                is_fixable=False,
                is_persistent=False,
                severity=ir.IssueSeverity.WARNING,
                translation_key="recommended_version",
                translation_placeholders={
                    "recommended_version": RECOMMENDED_VERSION,
                    "current_version": str(version),
                },
            )
    except Go2RtcClientError as err:
        if isinstance(err.__cause__, _RETRYABLE_ERRORS):
            raise ConfigEntryNotReady(
                f"Could not connect to go2rtc instance on {url}"
            ) from err
        _LOGGER.warning("Could not connect to go2rtc instance on %s (%s)", url, err)
        return False
    except Go2RtcVersionError as err:
        raise ConfigEntryNotReady(
            f"The go2rtc server version is not supported, {err}"
        ) from err
    except Exception as err:  # noqa: BLE001
        _LOGGER.warning("Could not connect to go2rtc instance on %s (%s)", url, err)
        return False

    provider = entry.runtime_data = WebRTCProvider(hass, url, session, client)
    await provider.initialize()
    entry.async_on_unload(async_register_webrtc_provider(hass, provider))
    return True


async def async_unload_entry(hass: HomeAssistant, entry: Go2RtcConfigEntry) -> bool:
    """Unload a go2rtc config entry."""
    await entry.runtime_data.teardown()
    return True


async def _get_binary(hass: HomeAssistant) -> str | None:
    """Return the binary path if found."""
    return await hass.async_add_executor_job(shutil.which, "go2rtc")


class WebRTCProvider(CameraWebRTCProvider):
    """WebRTC provider."""

    def __init__(
        self,
        hass: HomeAssistant,
        url: str,
        session: ClientSession,
        rest_client: Go2RtcRestClient,
    ) -> None:
        """Initialize the WebRTC provider."""
        self._hass = hass
        self._url = url
        self._session = session
        self._rest_client = rest_client
        self._sessions: dict[str, Go2RtcWsClient] = {}
        self._supported_schemes: set[str] = set()

    @property
    def domain(self) -> str:
        """Return the integration domain of the provider."""
        return DOMAIN

    async def initialize(self) -> None:
        """Initialize the provider."""
        self._supported_schemes = await self._rest_client.schemes.list()

    @callback
    def async_is_supported(self, stream_source: str) -> bool:
        """Return if this provider is supports the Camera as source."""
        return stream_source.partition(":")[0] in self._supported_schemes

    async def async_handle_async_webrtc_offer(
        self,
        camera: Camera,
        offer_sdp: str,
        session_id: str,
        send_message: WebRTCSendMessage,
    ) -> None:
        """Handle the WebRTC offer and return the answer via the provided callback."""
        try:
            await self._update_stream_source(camera)
        except HomeAssistantError as err:
            send_message(WebRTCError("go2rtc_webrtc_offer_failed", str(err)))
            return

        self._sessions[session_id] = ws_client = Go2RtcWsClient(
            self._session, self._url, source=camera.entity_id
        )

        @callback
        def on_messages(message: ReceiveMessages) -> None:
            """Handle messages."""
            value: WebRTCMessage
            match message:
                case WebRTCCandidate():
                    value = HAWebRTCCandidate(RTCIceCandidateInit(message.candidate))
                case WebRTCAnswer():
                    value = HAWebRTCAnswer(message.sdp)
                case WsError():
                    value = WebRTCError("go2rtc_webrtc_offer_failed", message.error)

            send_message(value)

        ws_client.subscribe(on_messages)
        config = camera.async_get_webrtc_client_configuration()
        await ws_client.send(WebRTCOffer(offer_sdp, config.configuration.ice_servers))

    async def async_on_webrtc_candidate(
        self, session_id: str, candidate: RTCIceCandidateInit
    ) -> None:
        """Handle the WebRTC candidate."""

        if ws_client := self._sessions.get(session_id):
            await ws_client.send(WebRTCCandidate(candidate.candidate))
        else:
            _LOGGER.debug("Unknown session %s. Ignoring candidate", session_id)

    @callback
    def async_close_session(self, session_id: str) -> None:
        """Close the session."""
        ws_client = self._sessions.pop(session_id)
        self._hass.async_create_task(ws_client.close())

    async def async_get_image(
        self,
        camera: Camera,
        width: int | None = None,
        height: int | None = None,
    ) -> bytes | None:
        """Get an image from the camera."""
        await self._update_stream_source(camera)
        return await self._rest_client.get_jpeg_snapshot(
            camera.entity_id, width, height
        )

    async def _update_stream_source(self, camera: Camera) -> None:
        """Update the stream source in go2rtc config if needed."""
        if not (stream_source := await camera.stream_source()):
            await self.teardown()
            raise HomeAssistantError("Camera has no stream source")

        if camera.platform.platform_name == "generic":
            # This is a workaround to use ffmpeg for generic cameras
            # A proper fix will be added in the future together with supporting multiple streams per camera
            stream_source = "ffmpeg:" + stream_source

        if not self.async_is_supported(stream_source):
            await self.teardown()
            raise HomeAssistantError("Stream source is not supported by go2rtc")

        camera_prefs = await get_dynamic_camera_stream_settings(
            self._hass, camera.entity_id
        )
        if camera_prefs.orientation is not Orientation.NO_TRANSFORM:
            # Camera orientation manually set by user
            if not stream_source.startswith(_FFMPEG):
                stream_source = _FFMPEG + ":" + stream_source
            stream_source += "#video=h264#audio=copy"
            match camera_prefs.orientation:
                case Orientation.MIRROR:
                    stream_source += "#raw=-vf hflip"
                case Orientation.ROTATE_180:
                    stream_source += "#rotate=180"
                case Orientation.FLIP:
                    stream_source += "#raw=-vf vflip"
                case Orientation.ROTATE_LEFT_AND_FLIP:
                    # Cannot use any filter when using raw one
                    stream_source += "#raw=-vf transpose=2,vflip"
                case Orientation.ROTATE_LEFT:
                    stream_source += "#rotate=-90"
                case Orientation.ROTATE_RIGHT_AND_FLIP:
                    # Cannot use any filter when using raw one
                    stream_source += "#raw=-vf transpose=1,vflip"
                case Orientation.ROTATE_RIGHT:
                    stream_source += "#rotate=90"

        streams = await self._rest_client.streams.list()

        if (stream := streams.get(camera.entity_id)) is None or not any(
            stream_source == producer.url for producer in stream.producers
        ):
            await self._rest_client.streams.add(
                camera.entity_id,
                [
                    stream_source,
                    # We are setting any ffmpeg rtsp related logs to debug
                    # Connection problems to the camera will be logged by the first stream
                    # Therefore setting it to debug will not hide any important logs
                    f"ffmpeg:{camera.entity_id}#audio=opus#query=log_level=debug",
                ],
            )

    async def teardown(self) -> None:
        """Tear down the provider."""
        for ws_client in self._sessions.values():
            await ws_client.close()
        self._sessions.clear()


@dataclass
class Go2RtcConfig:
    """Go2rtc configuration."""

    url: str
    session: ClientSession