Only accept images in allows MIME types

Author: outadoc

custom_components/immich/hub.py

@@ -11,6 +11,8 @@
 _HEADER_API_KEY = "x-api-key"
 _LOGGER = logging.getLogger(__name__)
 
+_ALLOWED_MIME_TYPES = ["image/png", "image/jpeg"]
+
 
 class ImmichHub:
     """Immich API hub."""
@@ -65,7 +67,7 @@ async def get_my_user_info(self) -> dict:
             _LOGGER.error("Error connecting to the API: %s", exception)
             raise CannotConnect from exception
 
-    async def download_asset(self, asset_id: str) -> bytes:
+    async def download_asset(self, asset_id: str) -> bytes | None:
         """Download the asset."""
         try:
             async with aiohttp.ClientSession() as session:
@@ -75,7 +77,13 @@ async def download_asset(self, asset_id: str) -> bytes:
                 async with session.get(url=url, headers=headers) as response:
                     if response.status != 200:
                         _LOGGER.error("Error from API: status=%d", response.status)
-                        raise ApiError()
+                        return None
+
+                    if response.content_type not in _ALLOWED_MIME_TYPES:
+                        _LOGGER.error(
+                            "MIME type is not supported: %s", response.content_type
+                        )
+                        return None
 
                     return await response.read()
         except aiohttp.ClientError as exception:

custom_components/immich/image.py

@@ -1,6 +1,7 @@
 """Image device for Immich integration."""
 from __future__ import annotations
 
+import asyncio
 from datetime import datetime, timedelta
 import logging
 import random
@@ -113,10 +114,20 @@ async def _get_next_asset_id(self) -> str | None:
 
     async def _load_and_cache_next_image(self) -> None:
         """Download and cache the image."""
-        asset_id = await self._get_next_asset_id()
+        asset_bytes = None
+
+        while not asset_bytes:
+            asset_id = await self._get_next_asset_id()
+
+            if not asset_id:
+                return
 
-        if asset_id:
             asset_bytes = await self.hub.download_asset(asset_id)
+
+            if not asset_bytes:
+                await asyncio.sleep(1)
+                continue
+
             self._current_image_bytes = asset_bytes
             self._attr_image_last_updated = datetime.now()
 

custom_components/immich/manifest.json

@@ -10,6 +10,6 @@
   "issue_tracker": "https://github.com/outadoc/immich-home-assistant/issues",
   "requirements": ["url-normalize==1.4.3"],
   "ssdp": [],
-  "version": "0.1.0",
+  "version": "0.2.0",
   "zeroconf": []
 }