Generic OTA providers

tests/conftest.py

@@ -104,7 +104,7 @@ def recursive_dict_merge(
     result = copy.deepcopy(obj)
 
     for key, update in updates.items():
-        if isinstance(update, dict):
+        if isinstance(update, dict) and key in result:
             result[key] = recursive_dict_merge(result[key], update)
         else:
             result[key] = update

tests/ota_providers/test_ota_provider_remote.py

@@ -0,0 +1,151 @@
+import hashlib
+import logging
+from unittest import mock
+
+import pytest
+
+import zigpy.ota
+import zigpy.ota.image
+from zigpy.ota.provider import RemoteImage, RemoteProvider
+
+from tests.async_mock import AsyncMock
+from tests.conftest import make_app
+
+MANUFACTURER_ID_1 = 0x1234
+MANUFACTURER_ID_2 = 0x5678
+IMAGE_TYPE = 0xABCD
+
+
+@pytest.fixture
+def provider():
+    p = RemoteProvider(
+        url="https://example.org/ota/",
+        manufacturer_ids=[MANUFACTURER_ID_1, MANUFACTURER_ID_2],
+    )
+    p.enable()
+    return p
+
+
+@pytest.fixture
+def ota_image():
+    img = zigpy.ota.image.OTAImage()
+    img.header = zigpy.ota.image.OTAImageHeader(
+        upgrade_file_id=zigpy.ota.image.OTAImageHeader.MAGIC_VALUE,
+        header_version=256,
+        header_length=56 + 2 + 2,
+        field_control=zigpy.ota.image.FieldControl.HARDWARE_VERSIONS_PRESENT,
+        manufacturer_id=MANUFACTURER_ID_2,
+        image_type=IMAGE_TYPE,
+        file_version=100,
+        stack_version=2,
+        header_string="This is a test header!",
+        image_size=56 + 2 + 4 + 4 + 2 + 2,
+        minimum_hardware_version=1,
+        maximum_hardware_version=3,
+    )
+    img.subelements = [zigpy.ota.image.SubElement(tag_id=0x0000, data=b"data")]
+
+    return img
+
+
+@pytest.fixture
+def image_json(ota_image):
+    return {
+        "binary_url": "https://example.org/ota/image1.ota",
+        "file_version": ota_image.header.file_version,
+        "image_type": ota_image.header.image_type,
+        "manufacturer_id": ota_image.header.manufacturer_id,
+        "changelog": "A changelog would go here.",
+        "checksum": f"sha3-256:{hashlib.sha3_256(ota_image.serialize()).hexdigest()}",
+        "min_hardware_version": ota_image.header.minimum_hardware_version,
+        "max_hardware_version": ota_image.header.maximum_hardware_version,
+        "min_current_file_version": 1,
+        "max_current_file_version": 99,
+    }
+
+
+@mock.patch("aiohttp.ClientSession.get")
+async def test_remote_image(mock_get, image_json, ota_image, provider, caplog):
+    image = RemoteImage.from_json(image_json)
+
+    assert image.key == zigpy.ota.image.ImageKey(
+        image.manufacturer_id,
+        image.image_type,
+    )
+
+    # Test unsuccessful download
+    rsp = mock_get.return_value.__aenter__.return_value
+    rsp.status = 404
+
+    with caplog.at_level(logging.WARNING):
+        await provider.initialize_provider({})
+
+    assert "Couldn't download" in caplog.text
+    caplog.clear()
+
+    # Test successful download
+    rsp.status = 200
+    rsp.json = AsyncMock(return_value=[image_json])
+    rsp.read = AsyncMock(return_value=ota_image.serialize())
+    await provider.initialize_provider({})
+
+    new_image = await provider.get_image(image.key)
+    assert new_image == ota_image
+
+
+@mock.patch("aiohttp.ClientSession.get")
+async def test_remote_image_bad_checksum(mock_get, image_json, ota_image, provider):
+    image = RemoteImage.from_json(image_json)
+
+    # Corrupt the checksum
+    image_json["checksum"] = f"sha3-256:{hashlib.sha3_256(b'').hexdigest()}"
+
+    # Test "successful" download
+    rsp = mock_get.return_value.__aenter__.return_value
+    rsp.status = 200
+    rsp.json = AsyncMock(return_value=[image_json])
+    rsp.read = AsyncMock(return_value=ota_image.serialize())
+    await provider.initialize_provider({})
+
+    # The image will fail to download
+    with pytest.raises(ValueError) as exc:
+        await provider.get_image(image.key)
+
+    assert "Image checksum is invalid" in str(exc.value)
+
+
+async def test_get_image_with_no_manufacturer_ids(ota_image, provider):
+    provider.manufacturer_ids = None
+
+    missing_key = zigpy.ota.image.ImageKey(
+        ota_image.header.manufacturer_id + 1,
+        ota_image.header.image_type + 1,
+    )
+
+    assert await provider.filter_get_image(missing_key) is False
+
+
+async def test_provider_initialization():
+    app = make_app(
+        {
+            "ota": {
+                "remote_providers": [
+                    {
+                        "url": "https://fw.zigbee.example.org/ota.json",
+                        "manufacturer_ids": [4660, 22136],
+                    },
+                    {
+                        "url": "https://fw.zigbee.example.org/ota-beta.json",
+                    },
+                ]
+            }
+        }
+    )
+
+    listeners, _ = zip(*app._ota._listeners.values())
+
+    assert listeners[0].url == "https://fw.zigbee.example.org/ota.json"
+    assert listeners[0].manufacturer_ids == [4660, 22136]
+
+    assert listeners[1].url == "https://fw.zigbee.example.org/ota-beta.json"
+    assert listeners[1].manufacturer_ids == []

zigpy/config/__init__.py

@@ -68,6 +68,9 @@
 CONF_OTA_SONOFF = "sonoff_provider"
 CONF_OTA_SONOFF_URL = "sonoff_update_url"
 CONF_OTA_THIRDREALITY = "thirdreality_provider"
+CONF_OTA_REMOTE_PROVIDERS = "remote_providers"
+CONF_OTA_PROVIDER_URL = "url"
+CONF_OTA_PROVIDER_MANUF_IDS = "manufacturer_ids"
 CONF_SOURCE_ROUTING = "source_routing"
 CONF_STARTUP_ENERGY_SCAN = "startup_energy_scan"
 CONF_TOPO_SCAN_PERIOD = "topology_scan_period"
@@ -105,6 +108,14 @@
         ),
     }
 )
+
+SCHEMA_OTA_PROVIDER = vol.Schema(
+    {
+        vol.Required(CONF_OTA_PROVIDER_URL): str,
+        vol.Optional(CONF_OTA_PROVIDER_MANUF_IDS, default=[]): [cv_hex],
+    }
+)
+
 SCHEMA_OTA = {
     vol.Optional(CONF_OTA_DIR, default=CONF_OTA_OTAU_DIR_DEFAULT): vol.Any(None, str),
     vol.Optional(CONF_OTA_IKEA, default=CONF_OTA_IKEA_DEFAULT): cv_boolean,
@@ -116,6 +127,7 @@
     vol.Optional(
         CONF_OTA_THIRDREALITY, default=CONF_OTA_THIRDREALITY_DEFAULT
     ): cv_boolean,
+    vol.Optional(CONF_OTA_REMOTE_PROVIDERS, default=[]): [SCHEMA_OTA_PROVIDER],
     # Deprecated keys
     vol.Optional(CONF_OTA_IKEA_URL): vol.All(
         cv_deprecated("The `ikea_update_url` key is deprecated and should be removed"),

zigpy/ota/__init__.py

@@ -12,6 +12,9 @@
     CONF_OTA_IKEA,
     CONF_OTA_INOVELLI,
     CONF_OTA_LEDVANCE,
+    CONF_OTA_PROVIDER_MANUF_IDS,
+    CONF_OTA_PROVIDER_URL,
+    CONF_OTA_REMOTE_PROVIDERS,
     CONF_OTA_SALUS,
     CONF_OTA_SONOFF,
     CONF_OTA_THIRDREALITY,
@@ -129,6 +132,14 @@ def __init__(self, app: ControllerApplicationType, *args, **kwargs) -> None:
         if ota_config[CONF_OTA_THIRDREALITY]:
             self.add_listener(zigpy.ota.provider.ThirdReality())
 
+        for provider_config in ota_config[CONF_OTA_REMOTE_PROVIDERS]:
+            self.add_listener(
+                zigpy.ota.provider.RemoteProvider(
+                    url=provider_config[CONF_OTA_PROVIDER_URL],
+                    manufacturer_ids=provider_config[CONF_OTA_PROVIDER_MANUF_IDS],
+                )
+            )
+
     async def initialize(self) -> None:
         await self.async_event("initialize_provider", self._app.config[CONF_OTA])
         self._not_initialized = False

zigpy/ota/provider.py

@@ -5,6 +5,7 @@
 import asyncio
 from collections import defaultdict
 import datetime
+import hashlib
 import io
 import logging
 import os
@@ -796,3 +797,107 @@
 
     async def filter_get_image(self, key: ImageKey) -> bool:
         return key.manufacturer_id not in self.MANUFACTURER_IDS
+
+
+@attr.s
+class RemoteImage:
+    binary_url = attr.ib()
+    file_version = attr.ib()
+    image_type = attr.ib()
+    manufacturer_id = attr.ib()
+    changelog = attr.ib()
+    checksum = attr.ib()
+
+    # Optional
+    min_hardware_version = attr.ib()
+    max_hardware_version = attr.ib()
+    min_current_file_version = attr.ib()
+    max_current_file_version = attr.ib()
+
+    @classmethod
+    def from_json(cls, obj: dict[str, typing.Any]) -> RemoteImage:
+        return cls(
+            binary_url=obj["binary_url"],
+            file_version=obj["file_version"],
+            image_type=obj["image_type"],
+            manufacturer_id=obj["manufacturer_id"],
+            changelog=obj["changelog"],
+            checksum=obj["checksum"],
+            min_hardware_version=obj.get("min_hardware_version"),
+            max_hardware_version=obj.get("max_hardware_version"),
+            min_current_file_version=obj.get("min_current_file_version"),
+            max_current_file_version=obj.get("max_current_file_version"),
+        )
+
+    @property
+    def key(self) -> ImageKey:
+        return ImageKey(self.manufacturer_id, self.image_type)
+
+    async def fetch_image(self) -> BaseOTAImage:
+        async with aiohttp.ClientSession() as req:
+            LOGGER.debug("Downloading %s for %s", self.binary_url, self.key)
+            async with req.get(self.binary_url) as rsp:
+                data = await rsp.read()
+
+        algorithm, checksum = self.checksum.split(":")
+        hasher = hashlib.new(algorithm)
+        await asyncio.get_running_loop().run_in_executor(None, hasher.update, data)
+
+        if hasher.hexdigest() != checksum:
+            raise ValueError(
+                f"Image checksum is invalid: expected {self.checksum},"
+                f" got {hasher.hexdigest()}"
+            )
+
+        ota_image, _ = parse_ota_image(data)
+
+        LOGGER.debug("Finished downloading %s", self)
+        return ota_image
+
+
+class RemoteProvider(Basic):
+    """Generic zigpy OTA URL provider."""
+
+    HEADERS = {"accept": "application/json"}
+
+    def __init__(self, url: str, manufacturer_ids: list[int] | None) -> None:
+        super().__init__()
+
+        self.url = url
+        self.manufacturer_ids = manufacturer_ids
+
+    async def initialize_provider(self, ota_config: dict) -> None:
+        self.info("OTA provider enabled")
+        await self.refresh_firmware_list()
+        self.enable()
+
+    async def refresh_firmware_list(self) -> None:
+        if self._locks[LOCK_REFRESH].locked():
+            return
+
+        async with self._locks[LOCK_REFRESH]:
+            async with aiohttp.ClientSession(headers=self.HEADERS) as req:
+                async with req.get(self.url) as rsp:
+                    if not (200 <= rsp.status <= 299):
+                        self.warning(
+                            "Couldn't download '%s': %s/%s",
+                            rsp.url,
+                            rsp.status,
+                            rsp.reason,
+                        )
+                        return
+                    fw_lst = await rsp.json()
+
+        self.debug("Finished downloading firmware update list")
+        self._cache.clear()
+        for obj in fw_lst:
+            img = RemoteImage.from_json(obj)
+            self._cache[img.key] = img
+
+        self.update_expiration()
+
+    async def filter_get_image(self, key: ImageKey) -> bool:
+        if not self.manufacturer_ids:
+            return False
+
+        return key.manufacturer_id not in self.manufacturer_ids