From 3c5056614ca2bea8b6718b24fd0d7ab1bc9ae6fa Mon Sep 17 00:00:00 2001 From: Sebastian Goscik Date: Sat, 18 Jan 2025 17:07:24 +0000 Subject: [PATCH] Monkey patch in experimental downloader --- .pre-commit-config.yaml | 1 + poetry.lock | 13 +- pyproject.toml | 4 + .../downloader_experimental.py | 2 - unifi_protect_backup/missing_event_checker.py | 2 +- unifi_protect_backup/uiprotect_patch.py | 135 ++++++++++++++++++ .../unifi_protect_backup_core.py | 8 ++ 7 files changed, 161 insertions(+), 4 deletions(-) create mode 100644 unifi_protect_backup/uiprotect_patch.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 840b1bb..041dc1e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -27,3 +27,4 @@ repos: - types-pytz - types-cryptography - types-python-dateutil + - types-aiofiles diff --git a/poetry.lock b/poetry.lock index a14eda7..96f3165 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1944,6 +1944,17 @@ rich = ">=10.11.0" shellingham = ">=1.3.0" typing-extensions = ">=3.7.4.3" +[[package]] +name = "types-aiofiles" +version = "24.1.0.20241221" +description = "Typing stubs for aiofiles" +optional = false +python-versions = ">=3.8" +files = [ + {file = "types_aiofiles-24.1.0.20241221-py3-none-any.whl", hash = "sha256:11d4e102af0627c02e8c1d17736caa3c39de1058bea37e2f4de6ef11a5b652ab"}, + {file = "types_aiofiles-24.1.0.20241221.tar.gz", hash = "sha256:c40f6c290b0af9e902f7f3fa91213cf5bb67f37086fb21dc0ff458253586ad55"}, +] + [[package]] name = "types-cryptography" version = "3.3.23.2" @@ -2181,4 +2192,4 @@ propcache = ">=0.2.0" [metadata] lock-version = "2.0" python-versions = ">=3.10.0,<4.0" -content-hash = "8a0fade4b4b3a1806c9c2fc2527f71b1b8d807b5b6abfe9dc49ee79364eba7d4" +content-hash = "2488bbbb25595c8758f278014438737e0f723e53b080244e36344fcd6081beea" diff --git a/pyproject.toml b/pyproject.toml index b9f126d..493f4d3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ types-python-dateutil = "^2.8.19.10" bump2version = "^1.0.1" pre-commit = "^2.12.0" ruff = "^0.5.7" +types-aiofiles = "^24.1.0.20241221" [tool.poetry.group.test] optional = true @@ -66,6 +67,9 @@ target-version = "py310" [tool.mypy] allow_redefinition=true +exclude = [ + 'unifi_protect_backup/uiprotect_patch.py' +] [build-system] requires = ["poetry-core>=1.0.0"] diff --git a/unifi_protect_backup/downloader_experimental.py b/unifi_protect_backup/downloader_experimental.py index d4a7a3a..e336ba0 100644 --- a/unifi_protect_backup/downloader_experimental.py +++ b/unifi_protect_backup/downloader_experimental.py @@ -85,8 +85,6 @@ def __init__( else: self._has_ffprobe = False - raise RuntimeError("The `uiprotect` library is currently missing the features for this to work.") - async def start(self): """Main loop.""" self.logger.info("Starting Downloader") diff --git a/unifi_protect_backup/missing_event_checker.py b/unifi_protect_backup/missing_event_checker.py index 7572ffc..3d1286a 100644 --- a/unifi_protect_backup/missing_event_checker.py +++ b/unifi_protect_backup/missing_event_checker.py @@ -87,7 +87,7 @@ async def _get_missing_events(self) -> AsyncIterator[Event]: break # No completed events to process # Next chunks start time should be the end of the oldest complete event in the current chunk - start_time = max([event.end for event in unifi_events.values()]) + start_time = max([event.end for event in unifi_events.values() if event.end is not None]) # Get list of events that have been backed up from the database diff --git a/unifi_protect_backup/uiprotect_patch.py b/unifi_protect_backup/uiprotect_patch.py new file mode 100644 index 0000000..278f977 --- /dev/null +++ b/unifi_protect_backup/uiprotect_patch.py @@ -0,0 +1,135 @@ +import enum +from datetime import datetime +from pathlib import Path +from typing import Any, Optional + +import aiofiles + +from uiprotect.data import Version +from uiprotect.exceptions import BadRequest +from uiprotect.utils import to_js_time + + +# First, let's add the new VideoExportType enum +class VideoExportType(str, enum.Enum): + TIMELAPSE = "timelapse" + ROTATING = "rotating" + + +def monkey_patch_experimental_downloader(): + from uiprotect.api import ProtectApiClient + + # Add the version constant + ProtectApiClient.NEW_DOWNLOAD_VERSION = Version("4.0.0") # You'll need to import Version from uiprotect + + async def _validate_channel_id(self, camera_id: str, channel_index: int) -> None: + if self._bootstrap is None: + await self.update() + try: + camera = self._bootstrap.cameras[camera_id] + camera.channels[channel_index] + except (IndexError, AttributeError, KeyError) as e: + raise BadRequest(f"Invalid input: {e}") from e + + async def prepare_camera_video( + self, + camera_id: str, + start: datetime, + end: datetime, + channel_index: int = 0, + validate_channel_id: bool = True, + fps: Optional[int] = None, + filename: Optional[str] = None, + ) -> Optional[dict[str, Any]]: + if self.bootstrap.nvr.version < self.NEW_DOWNLOAD_VERSION: + raise ValueError("This method is only support from Unifi Protect version >= 4.0.0.") + + if validate_channel_id: + await self._validate_channel_id(camera_id, channel_index) + + params = { + "camera": camera_id, + "start": to_js_time(start), + "end": to_js_time(end), + } + + if channel_index == 3: + params.update({"lens": 2}) + else: + params.update({"channel": channel_index}) + + if fps is not None and fps > 0: + params["fps"] = fps + params["type"] = VideoExportType.TIMELAPSE.value + else: + params["type"] = VideoExportType.ROTATING.value + + if not filename: + start_str = start.strftime("%m-%d-%Y, %H.%M.%S %Z") + end_str = end.strftime("%m-%d-%Y, %H.%M.%S %Z") + filename = f"{camera_id} {start_str} - {end_str}.mp4" + + params["filename"] = filename + + return await self.api_request( + "video/prepare", + params=params, + raise_exception=True, + ) + + async def download_camera_video( + self, + camera_id: str, + filename: str, + output_file: Optional[Path] = None, + iterator_callback: Optional[callable] = None, + progress_callback: Optional[callable] = None, + chunk_size: int = 65536, + ) -> Optional[bytes]: + if self.bootstrap.nvr.version < self.NEW_DOWNLOAD_VERSION: + raise ValueError("This method is only support from Unifi Protect version >= 4.0.0.") + + params = { + "camera": camera_id, + "filename": filename, + } + + if iterator_callback is None and progress_callback is None and output_file is None: + return await self.api_request_raw( + "video/download", + params=params, + raise_exception=False, + ) + + r = await self.request( + "get", + f"{self.api_path}video/download", + auto_close=False, + timeout=0, + params=params, + ) + + if output_file is not None: + async with aiofiles.open(output_file, "wb") as output: + + async def callback(total: int, chunk: Optional[bytes]) -> None: + if iterator_callback is not None: + await iterator_callback(total, chunk) + if chunk is not None: + await output.write(chunk) + + await self._stream_response(r, chunk_size, callback, progress_callback) + else: + await self._stream_response( + r, + chunk_size, + iterator_callback, + progress_callback, + ) + r.close() + return None + + # Patch the methods into the class + ProtectApiClient._validate_channel_id = _validate_channel_id + ProtectApiClient.prepare_camera_video = prepare_camera_video + ProtectApiClient.download_camera_video = download_camera_video diff --git a/unifi_protect_backup/unifi_protect_backup_core.py b/unifi_protect_backup/unifi_protect_backup_core.py index 4cdbb5b..b31508c 100644 --- a/unifi_protect_backup/unifi_protect_backup_core.py +++ b/unifi_protect_backup/unifi_protect_backup_core.py @@ -29,11 +29,19 @@ setup_logging, ) +from unifi_protect_backup.uiprotect_patch import monkey_patch_experimental_downloader + logger = logging.getLogger(__name__) # TODO: https://github.com/cjrh/aiorun#id6 (smart shield) +# We have been waiting for a long time for this PR to get merged +# https://github.com/uilibs/uiprotect/pull/249 +# Since it has not progressed, we will for now patch in the functionality ourselves +monkey_patch_experimental_downloader() + + async def create_database(path: str): """Creates sqlite database and creates the events abd backups tables.""" db = await aiosqlite.connect(path)