uiprotect 1.20.0__tar.gz → 2.0.0__tar.gz

{uiprotect-1.20.0 → uiprotect-2.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: uiprotect
-Version: 1.20.0
+Version: 2.0.0
 Summary: Python API for Unifi Protect (Unofficial)
 Home-page: https://github.com/uilibs/uiprotect
 License: MIT

{uiprotect-1.20.0 → uiprotect-2.0.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "uiprotect"
-version = "1.20.0"
+version = "2.0.0"
 description = "Python API for Unifi Protect (Unofficial)"
 authors = ["UI Protect Maintainers <ui@koston.org>"]
 license = "MIT"

{uiprotect-1.20.0 → uiprotect-2.0.0}/src/uiprotect/api.py

@@ -11,7 +11,7 @@ import sys
 import time
 from collections.abc import Callable
 from datetime import datetime, timedelta
-from functools import cached_property
+from functools import cached_property, partial
 from http.cookies import Morsel, SimpleCookie
 from ipaddress import IPv4Address, IPv6Address
 from pathlib import Path
@@ -203,6 +203,7 @@ class BaseApiClient:
         self._verify_ssl = verify_ssl
         self._ws_timeout = ws_timeout
         self._loaded_session = False
+        self._update_task: asyncio.Task[Bootstrap | None] | None = None
 
         self.config_dir = config_dir or (Path(user_config_dir()) / "ufp")
         self.cache_dir = cache_dir or (Path(user_cache_dir()) / "ufp_cache")
@@ -253,37 +254,56 @@ class BaseApiClient:
 
         return self._session
 
-    async def get_websocket(self) -> Websocket:
-        """Gets or creates current Websocket."""
-
-        async def _auth(force: bool) -> dict[str, str] | None:
-            if force:
-                if self._session is not None:
-                    self._session.cookie_jar.clear()
-                self.set_header("cookie", None)
-                self.set_header("x-csrf-token", None)
+    async def _auth_websocket(self, force: bool) -> dict[str, str] | None:
+        """Authenticate for Websocket."""
+        if force:
+            if self._session is not None:
+                self._session.cookie_jar.clear()
+            self.set_header("cookie", None)
+            self.set_header("x-csrf-token", None)
+            self._is_authenticated = False
 
-            await self.ensure_authenticated()
-            return self.headers
+        await self.ensure_authenticated()
+        return self.headers
 
+    def _get_websocket(self) -> Websocket:
+        """Gets or creates current Websocket."""
         if self._websocket is None:
             self._websocket = Websocket(
                 self.get_websocket_url,
-                _auth,
+                self._auth_websocket,
+                self._update_bootstrap_soon,
+                self.get_session,
+                self._process_ws_message,
                 verify=self._verify_ssl,
                 timeout=self._ws_timeout,
             )
-            self._websocket.subscribe(self._process_ws_message)
-
         return self._websocket
 
+    def _update_bootstrap_soon(self) -> None:
+        """Update bootstrap soon."""
+        _LOGGER.debug("Updating bootstrap soon")
+        # Force the next bootstrap update
+        # since the lastUpdateId is not valid anymore
+        if self._update_task and not self._update_task.done():
+            return
+        self._update_task = asyncio.create_task(self.update(force=True))
+
     async def close_session(self) -> None:
-        """Closing and delets client session"""
+        """Closing and deletes client session"""
+        await self._cancel_update_task()
         if self._session is not None:
             await self._session.close()
             self._session = None
             self._loaded_session = False
 
+    async def _cancel_update_task(self) -> None:
+        if self._update_task:
+            self._update_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._update_task
+            self._update_task = None
+
     def set_header(self, key: str, value: str | None) -> None:
         """Set header."""
         self.headers = self.headers or {}
@@ -413,8 +433,13 @@ class BaseApiClient:
         )
 
         if data is not None:
-            json_data: list[Any] | dict[str, Any] = orjson.loads(data)
-            return json_data
+            json_data: list[Any] | dict[str, Any]
+            try:
+                json_data = orjson.loads(data)
+                return json_data
+            except orjson.JSONDecodeError as ex:
+                _LOGGER.error("Could not decode JSON from %s", url)
+                raise NvrError(f"Could not decode JSON from {url}") from ex
         return None
 
     async def api_request_obj(
@@ -620,31 +645,17 @@ class BaseApiClient:
 
         return token_expires_at >= max_expire_time
 
-    async def async_connect_ws(self, force: bool) -> None:
-        """Connect to Websocket."""
-        if force and self._websocket is not None:
-            await self._websocket.disconnect()
-            self._websocket = None
-
-        websocket = await self.get_websocket()
-        if not websocket.is_connected:
-            self._last_ws_status = False
-            with contextlib.suppress(
-                TimeoutError,
-                asyncio.TimeoutError,
-                asyncio.CancelledError,
-            ):
-                await websocket.connect()
-
     def get_websocket_url(self) -> str:
         """Get Websocket URL."""
         return self.ws_url
 
     async def async_disconnect_ws(self) -> None:
         """Disconnect from Websocket."""
-        if self._websocket is None:
-            return
-        await self._websocket.disconnect()
+        if self._websocket:
+            websocket = self._get_websocket()
+            websocket.stop()
+            await websocket.wait_closed()
+            self._websocket = None
 
     def check_ws(self) -> bool:
         """Checks current state of Websocket."""
@@ -668,6 +679,9 @@ class BaseApiClient:
     def _get_last_update_id(self) -> str | None:
         raise NotImplementedError
 
+    async def update(self, force: bool = False) -> Bootstrap | None:
+        raise NotImplementedError
+
 
 class ProtectApiClient(BaseApiClient):
     """
@@ -748,6 +762,7 @@ class ProtectApiClient(BaseApiClient):
         self._ignore_stats = ignore_stats
         self._ws_subscriptions = []
         self.ignore_unadopted = ignore_unadopted
+        self._update_lock = asyncio.Lock()
 
         if override_connection_host:
             self._connection_host = ip_from_host(self._host)
@@ -788,41 +803,37 @@ class ProtectApiClient(BaseApiClient):
 
         You can use the various other `get_` methods if you need one off data from UFP
         """
-        now = time.monotonic()
-        now_dt = utc_now()
-        max_event_dt = now_dt - timedelta(hours=1)
-        if force:
-            self._last_update = NEVER_RAN
-            self._last_update_dt = max_event_dt
-
-        bootstrap_updated = False
-        if self._bootstrap is None or now - self._last_update > DEVICE_UPDATE_INTERVAL:
-            bootstrap_updated = True
-            self._bootstrap = await self.get_bootstrap()
-            self.__dict__.pop("bootstrap", None)
-            self._last_update = now
-            self._last_update_dt = now_dt
+        async with self._update_lock:
+            now = time.monotonic()
+            if force:
+                self._last_update = NEVER_RAN
 
-        await self.async_connect_ws(force)
-        if self.check_ws():
-            # If the websocket is connected/connecting
-            # we do not need to get events
-            _LOGGER.debug("Skipping update since websocket is active")
-            return None
+            bootstrap_updated = False
+            if (
+                self._bootstrap is None
+                or now - self._last_update > DEVICE_UPDATE_INTERVAL
+            ):
+                bootstrap_updated = True
+                self._bootstrap = await self.get_bootstrap()
+                self.__dict__.pop("bootstrap", None)
+                self._last_update = now
 
-        if bootstrap_updated:
-            return None
+            if bootstrap_updated:
+                return None
+            self._last_update = now
+            return self._bootstrap
 
+    async def poll_events(self) -> None:
+        """Poll for events."""
+        now_dt = utc_now()
+        max_event_dt = now_dt - timedelta(hours=1)
         events = await self.get_events(
             start=self._last_update_dt or max_event_dt,
             end=now_dt,
         )
         for event in events:
             self.bootstrap.process_event(event)
-
-        self._last_update = now
         self._last_update_dt = now_dt
-        return self._bootstrap
 
     def emit_message(self, msg: WSSubscriptionMessage) -> None:
         """Emit message to all subscriptions."""
@@ -1108,13 +1119,20 @@ class ProtectApiClient(BaseApiClient):
 
         Returns a callback that will unsubscribe.
         """
-
-        def _unsub_ws_callback() -> None:
-            self._ws_subscriptions.remove(ws_callback)
-
         _LOGGER.debug("Adding subscription: %s", ws_callback)
         self._ws_subscriptions.append(ws_callback)
-        return _unsub_ws_callback
+        self._get_websocket().start()
+        return partial(self._unsubscribe_websocket, ws_callback)
+
+    def _unsubscribe_websocket(
+        self,
+        ws_callback: Callable[[WSSubscriptionMessage], None],
+    ) -> None:
+        """Unsubscribe to websocket events."""
+        _LOGGER.debug("Removing subscription: %s", ws_callback)
+        self._ws_subscriptions.remove(ws_callback)
+        if not self._ws_subscriptions:
+            self._get_websocket().stop()
 
     async def get_bootstrap(self) -> Bootstrap:
         """

{uiprotect-1.20.0 → uiprotect-2.0.0}/src/uiprotect/cli/__init__.py

@@ -201,6 +201,7 @@ def shell(ctx: typer.Context) -> None:
 
     async def wait_forever() -> None:
         await protect.update()
+        protect.subscribe_websocket(lambda _: None)
         while True:
             await asyncio.sleep(10)
             await protect.update()
@@ -262,12 +263,16 @@ def profile_ws(
 
     async def callback() -> None:
         await protect.update()
+        unsub = protect.subscribe_websocket(lambda _: None)
         await profile_ws_job(
             protect,
             wait_time,
             output_path=output_path,
             ws_progress=_progress_bar,
         )
+        unsub()
+        await protect.async_disconnect_ws()
+        await protect.close_session()
 
     _setup_logger()
 

{uiprotect-1.20.0 → uiprotect-2.0.0}/src/uiprotect/test_util/__init__.py

@@ -103,8 +103,10 @@ class SampleDataGenerator:
     async def async_generate(self, close_session: bool = True) -> None:
         self.log(f"Output folder: {self.output_folder}")
         self.output_folder.mkdir(parents=True, exist_ok=True)
-        websocket = await self.client.get_websocket()
-        websocket.subscribe(self._handle_ws_message)
+        websocket = self.client._get_websocket()
+        websocket.start()
+        self.log("Websocket started...")
+        websocket._subscription = self._handle_ws_message
 
         self.log("Updating devices...")
         await self.client.update()
@@ -131,8 +133,10 @@ class SampleDataGenerator:
             "chime": len(bootstrap["chimes"]),
         }
 
+        self.log("Generating event data...")
         motion_event, smart_detection = await self.generate_event_data()
         await self.generate_device_data(motion_event, smart_detection)
+        self.log("Recording websocket events...")
         await self.record_ws_events()
 
         if close_session:

uiprotect-2.0.0/src/uiprotect/websocket.py

@@ -0,0 +1,183 @@
+"""UniFi Protect Websockets."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+from collections.abc import Awaitable, Callable, Coroutine
+from http import HTTPStatus
+from typing import Any, Optional
+
+from aiohttp import (
+    ClientError,
+    ClientSession,
+    ClientWebSocketResponse,
+    WSMessage,
+    WSMsgType,
+    WSServerHandshakeError,
+)
+
+_LOGGER = logging.getLogger(__name__)
+AuthCallbackType = Callable[..., Coroutine[Any, Any, Optional[dict[str, str]]]]
+GetSessionCallbackType = Callable[[], Awaitable[ClientSession]]
+UpdateBootstrapCallbackType = Callable[[], None]
+_CLOSE_MESSAGE_TYPES = {WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED}
+
+
+class Websocket:
+    """UniFi Protect Websocket manager."""
+
+    _running = False
+    _headers: dict[str, str] | None = None
+    _websocket_loop_task: asyncio.Task[None] | None = None
+    _stop_task: asyncio.Task[None] | None = None
+    _ws_connection: ClientWebSocketResponse | None = None
+
+    def __init__(
+        self,
+        get_url: Callable[[], str],
+        auth_callback: AuthCallbackType,
+        update_bootstrap_callback: UpdateBootstrapCallbackType,
+        get_session: GetSessionCallbackType,
+        subscription: Callable[[WSMessage], None],
+        *,
+        timeout: float = 30.0,
+        backoff: int = 10,
+        verify: bool = True,
+    ) -> None:
+        """Init Websocket."""
+        self.get_url = get_url
+        self.timeout = timeout
+        self.backoff = backoff
+        self.verify = verify
+        self._get_session = get_session
+        self._auth = auth_callback
+        self._update_bootstrap_callback = update_bootstrap_callback
+        self._connect_lock = asyncio.Lock()
+        self._subscription = subscription
+        self._last_ws_connect_ok = False
+
+    @property
+    def is_connected(self) -> bool:
+        """Return if the websocket is connected."""
+        return self._ws_connection is not None and not self._ws_connection.closed
+
+    async def _websocket_reconnect_loop(self) -> None:
+        """Reconnect loop for websocket."""
+        await self.wait_closed()
+        backoff = self.backoff
+
+        while True:
+            try:
+                await self._websocket_loop()
+            except ClientError:
+                _LOGGER.debug("Error in websocket reconnect loop, backoff: %s", backoff)
+            except Exception:
+                _LOGGER.debug(
+                    "Error in websocket reconnect loop, backoff: %s",
+                    backoff,
+                    exc_info=True,
+                )
+
+            if self._running is False:
+                break
+            await asyncio.sleep(self.backoff)
+
+    async def _websocket_loop(self) -> None:
+        url = self.get_url()
+        _LOGGER.debug("Connecting WS to %s", url)
+        self._headers = await self._auth(False)
+        ssl = None if self.verify else False
+        msg: WSMessage | None = None
+        seen_non_close_message = False
+        # catch any and all errors for Websocket so we can clean up correctly
+        try:
+            session = await self._get_session()
+            self._ws_connection = await session.ws_connect(
+                url, ssl=ssl, headers=self._headers, timeout=self.timeout
+            )
+            self._last_ws_connect_ok = True
+            while True:
+                msg = await self._ws_connection.receive(self.timeout)
+                msg_type = msg.type
+                if msg_type is WSMsgType.ERROR:
+                    _LOGGER.exception("Error from Websocket: %s", msg.data)
+                    break
+                elif msg_type in _CLOSE_MESSAGE_TYPES:
+                    _LOGGER.debug("Websocket closed: %s", msg)
+                    break
+
+                seen_non_close_message = True
+                try:
+                    self._subscription(msg)
+                except Exception:
+                    _LOGGER.exception("Error processing websocket message")
+        except asyncio.TimeoutError:
+            _LOGGER.debug("Websocket timeout: %s", url)
+        except WSServerHandshakeError as ex:
+            level = logging.ERROR if self._last_ws_connect_ok else logging.DEBUG
+            self._last_ws_connect_ok = False
+            if ex.status == HTTPStatus.UNAUTHORIZED.value:
+                _LOGGER.log(level, "Websocket authentication error: %s", url)
+                self._headers = await self._auth(True)
+            else:
+                _LOGGER.log(level, "Websocket handshake error: %s", url, exc_info=True)
+            raise
+        except ClientError:
+            level = logging.ERROR if self._last_ws_connect_ok else logging.DEBUG
+            self._last_ws_connect_ok = False
+            _LOGGER.log(level, "Websocket disconnect error: %s", url, exc_info=True)
+            raise
+        finally:
+            if (
+                msg is not None
+                and msg.type is WSMsgType.CLOSE
+                # If it closes right away or lastUpdateId is in the extra
+                # its an indication that we should update the bootstrap
+                # since lastUpdateId is invalid
+                and (
+                    not seen_non_close_message
+                    or (msg.extra and "lastUpdateId" in msg.extra)
+                )
+            ):
+                self._update_bootstrap_callback()
+            _LOGGER.debug("Websocket disconnected: last message: %s", msg)
+            if self._ws_connection is not None and not self._ws_connection.closed:
+                await self._ws_connection.close()
+            self._ws_connection = None
+
+    def start(self) -> None:
+        """Start the websocket."""
+        if self._running:
+            return
+        self._running = True
+        self._websocket_loop_task = asyncio.create_task(
+            self._websocket_reconnect_loop()
+        )
+
+    def stop(self) -> None:
+        """Disconnect the websocket."""
+        _LOGGER.debug("Disconnecting websocket...")
+        if not self._running:
+            return
+        if self._websocket_loop_task:
+            self._websocket_loop_task.cancel()
+        self._running = False
+        self._stop_task = asyncio.create_task(self._stop())
+
+    async def wait_closed(self) -> None:
+        """Wait for the websocket to close."""
+        if self._stop_task:
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._stop_task
+
+    async def _stop(self) -> None:
+        """Stop the websocket."""
+        if self._ws_connection:
+            await self._ws_connection.close()
+            self._ws_connection = None
+        if self._websocket_loop_task:
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._websocket_loop_task
+            self._websocket_loop_task = None

uiprotect-1.20.0/src/uiprotect/websocket.py

@@ -1,226 +0,0 @@
-"""UniFi Protect Websockets."""
-
-from __future__ import annotations
-
-import asyncio
-import logging
-import time
-from collections.abc import Callable, Coroutine
-from typing import Any, Optional
-
-from aiohttp import (
-    ClientError,
-    ClientSession,
-    ClientWebSocketResponse,
-    WSMessage,
-    WSMsgType,
-)
-
-from .utils import asyncio_timeout
-
-_LOGGER = logging.getLogger(__name__)
-CALLBACK_TYPE = Callable[..., Coroutine[Any, Any, Optional[dict[str, str]]]]
-RECENT_FAILURE_CUT_OFF = 30
-RECENT_FAILURE_THRESHOLD = 2
-
-
-class Websocket:
-    """UniFi Protect Websocket manager."""
-
-    url: str
-    verify: bool
-    timeout_interval: int
-    backoff: int
-    _auth: CALLBACK_TYPE
-    _timeout: float
-    _ws_subscriptions: list[Callable[[WSMessage], None]]
-    _connect_lock: asyncio.Lock
-
-    _headers: dict[str, str] | None = None
-    _websocket_loop_task: asyncio.Task[None] | None = None
-    _timer_task: asyncio.Task[None] | None = None
-    _ws_connection: ClientWebSocketResponse | None = None
-    _last_connect: float = -1000
-    _recent_failures: int = 0
-
-    def __init__(
-        self,
-        get_url: Callable[[], str],
-        auth_callback: CALLBACK_TYPE,
-        *,
-        timeout: int = 30,
-        backoff: int = 10,
-        verify: bool = True,
-    ) -> None:
-        """Init Websocket."""
-        self.get_url = get_url
-        self.timeout_interval = timeout
-        self.backoff = backoff
-        self.verify = verify
-        self._auth = auth_callback
-        self._timeout = time.monotonic()
-        self._ws_subscriptions = []
-        self._connect_lock = asyncio.Lock()
-
-    @property
-    def is_connected(self) -> bool:
-        """Check if Websocket connected."""
-        return self._ws_connection is not None
-
-    def _get_session(self) -> ClientSession:
-        # for testing, to make easier to mock
-        return ClientSession()
-
-    def _process_message(self, msg: WSMessage) -> bool:
-        if msg.type == WSMsgType.ERROR:
-            _LOGGER.exception("Error from Websocket: %s", msg.data)
-            return False
-
-        for sub in self._ws_subscriptions:
-            try:
-                sub(msg)
-            except Exception:
-                _LOGGER.exception("Error processing websocket message")
-
-        return True
-
-    async def _websocket_loop(self, start_event: asyncio.Event) -> None:
-        url = self.get_url()
-        _LOGGER.debug("Connecting WS to %s", url)
-        self._headers = await self._auth(self._should_reset_auth)
-
-        session = self._get_session()
-        # catch any and all errors for Websocket so we can clean up correctly
-        try:
-            self._ws_connection = await session.ws_connect(
-                url,
-                ssl=None if self.verify else False,
-                headers=self._headers,
-            )
-            start_event.set()
-
-            self._reset_timeout()
-            async for msg in self._ws_connection:
-                if not self._process_message(msg):
-                    break
-                self._reset_timeout()
-        except ClientError:
-            _LOGGER.exception("Websocket disconnect error: %s", url)
-        finally:
-            _LOGGER.debug("Websocket disconnected")
-            self._increase_failure()
-            self._cancel_timeout()
-            if self._ws_connection is not None and not self._ws_connection.closed:
-                await self._ws_connection.close()
-            if not session.closed:
-                await session.close()
-            self._ws_connection = None
-            # make sure event does not timeout
-            start_event.set()
-
-    @property
-    def has_recent_connect(self) -> bool:
-        """Check if Websocket has recent connection."""
-        return time.monotonic() - RECENT_FAILURE_CUT_OFF <= self._last_connect
-
-    @property
-    def _should_reset_auth(self) -> bool:
-        if self.has_recent_connect:
-            if self._recent_failures > RECENT_FAILURE_THRESHOLD:
-                return True
-        else:
-            self._recent_failures = 0
-        return False
-
-    def _increase_failure(self) -> None:
-        if self.has_recent_connect:
-            self._recent_failures += 1
-        else:
-            self._recent_failures = 1
-
-    async def _do_timeout(self) -> bool:
-        _LOGGER.debug("WS timed out")
-        return await self.reconnect()
-
-    async def _timeout_loop(self) -> None:
-        while True:
-            now = time.monotonic()
-            if now > self._timeout:
-                _LOGGER.debug("WS timed out")
-                if not await self.reconnect():
-                    _LOGGER.debug("WS could not reconnect")
-                    continue
-            sleep_time = self._timeout - now
-            _LOGGER.debug("WS Timeout loop sleep %s", sleep_time)
-            await asyncio.sleep(sleep_time)
-
-    def _reset_timeout(self) -> None:
-        self._timeout = time.monotonic() + self.timeout_interval
-
-        if self._timer_task is None:
-            self._timer_task = asyncio.create_task(self._timeout_loop())
-
-    def _cancel_timeout(self) -> None:
-        if self._timer_task:
-            self._timer_task.cancel()
-
-    async def connect(self) -> bool:
-        """Connect the websocket."""
-        if self._connect_lock.locked():
-            _LOGGER.debug("Another connect is already happening")
-            return False
-        try:
-            async with asyncio_timeout(0.1):
-                await self._connect_lock.acquire()
-        except (TimeoutError, asyncio.TimeoutError, asyncio.CancelledError):
-            _LOGGER.debug("Failed to get connection lock")
-
-        start_event = asyncio.Event()
-        _LOGGER.debug("Scheduling WS connect...")
-        self._websocket_loop_task = asyncio.create_task(
-            self._websocket_loop(start_event),
-        )
-
-        try:
-            async with asyncio_timeout(self.timeout_interval):
-                await start_event.wait()
-        except (TimeoutError, asyncio.TimeoutError, asyncio.CancelledError):
-            _LOGGER.warning("Timed out while waiting for Websocket to connect")
-            await self.disconnect()
-
-        self._connect_lock.release()
-        if self._ws_connection is None:
-            _LOGGER.debug("Failed to connect to Websocket")
-            return False
-        _LOGGER.debug("Connected to Websocket successfully")
-        self._last_connect = time.monotonic()
-        return True
-
-    async def disconnect(self) -> None:
-        """Disconnect the websocket."""
-        _LOGGER.debug("Disconnecting websocket...")
-        if self._ws_connection is None:
-            return
-        await self._ws_connection.close()
-        self._ws_connection = None
-
-    async def reconnect(self) -> bool:
-        """Reconnect the websocket."""
-        _LOGGER.debug("Reconnecting websocket...")
-        await self.disconnect()
-        await asyncio.sleep(self.backoff)
-        return await self.connect()
-
-    def subscribe(self, ws_callback: Callable[[WSMessage], None]) -> Callable[[], None]:
-        """
-        Subscribe to raw websocket messages.
-
-        Returns a callback that will unsubscribe.
-        """
-
-        def _unsub_ws_callback() -> None:
-            self._ws_subscriptions.remove(ws_callback)
-
-        _LOGGER.debug("Adding subscription: %s", ws_callback)
-        self._ws_subscriptions.append(ws_callback)
-        return _unsub_ws_callback