"""Implementation of firmware module.""" from __future__ import annotations import asyncio import logging from asyncio import timeout as asyncio_timeout from collections.abc import Callable, Coroutine from dataclasses import dataclass, field from datetime import date from typing import TYPE_CHECKING, Annotated from mashumaro import DataClassDictMixin, field_options from mashumaro.types import Alias from ...exceptions import KasaException from ...feature import Feature from ..smartmodule import SmartModule, allow_update_after if TYPE_CHECKING: from ..smartdevice import SmartDevice _LOGGER = logging.getLogger(__name__) @dataclass class DownloadState(DataClassDictMixin): """Download state.""" # Example: # {'status': 0, 'download_progress': 0, 'reboot_time': 5, # 'upgrade_time': 5, 'auto_upgrade': False} status: int progress: Annotated[int, Alias("download_progress")] reboot_time: int upgrade_time: int auto_upgrade: bool @dataclass class UpdateInfo(DataClassDictMixin): """Update info status object.""" status: Annotated[int, Alias("type")] needs_upgrade: Annotated[bool, Alias("need_to_upgrade")] version: Annotated[str | None, Alias("fw_ver")] = None release_date: date | None = field( default=None, metadata=field_options( deserialize=lambda x: date.fromisoformat(x) if x else None ), ) release_notes: Annotated[str | None, Alias("release_note")] = None fw_size: int | None = None oem_id: str | None = None @property def update_available(self) -> bool: """Return True if update available.""" return self.status != 0 class Firmware(SmartModule): """Implementation of firmware module.""" REQUIRED_COMPONENT = "firmware" MINIMUM_UPDATE_INTERVAL_SECS = 60 * 60 * 24 def __init__(self, device: SmartDevice, module: str) -> None: super().__init__(device, module) self._firmware_update_info: UpdateInfo | None = None def _initialize_features(self) -> None: """Initialize features.""" device = self._device if self.supported_version > 1: self._add_feature( Feature( device, id="auto_update_enabled", name="Auto update enabled", container=self, attribute_getter="auto_update_enabled", attribute_setter="set_auto_update_enabled", type=Feature.Type.Switch, ) ) self._add_feature( Feature( device, id="update_available", name="Update available", container=self, attribute_getter="update_available", type=Feature.Type.BinarySensor, category=Feature.Category.Info, ) ) self._add_feature( Feature( device, id="current_firmware_version", name="Current firmware version", container=self, attribute_getter="current_firmware", category=Feature.Category.Debug, type=Feature.Type.Sensor, ) ) self._add_feature( Feature( device, id="available_firmware_version", name="Available firmware version", container=self, attribute_getter="latest_firmware", category=Feature.Category.Debug, type=Feature.Type.Sensor, ) ) self._add_feature( Feature( device, id="check_latest_firmware", name="Check latest firmware", container=self, attribute_setter="check_latest_firmware", category=Feature.Category.Info, type=Feature.Type.Action, ) ) def query(self) -> dict: """Query to execute during the update cycle.""" if self.supported_version > 1: return {"get_auto_update_info": None} return {} async def check_latest_firmware(self) -> UpdateInfo | None: """Check for the latest firmware for the device.""" try: fw = await self.call("get_latest_fw") self._firmware_update_info = UpdateInfo.from_dict(fw["get_latest_fw"]) return self._firmware_update_info except Exception: _LOGGER.exception("Error getting latest firmware for %s:", self._device) self._firmware_update_info = None return None @property def current_firmware(self) -> str: """Return the current firmware version.""" return self._device.hw_info["sw_ver"] @property def latest_firmware(self) -> str | None: """Return the latest firmware version.""" if not self._firmware_update_info: return None return self._firmware_update_info.version @property def firmware_update_info(self) -> UpdateInfo | None: """Return latest firmware information.""" return self._firmware_update_info @property def update_available(self) -> bool | None: """Return True if update is available.""" if not self._device.is_cloud_connected or not self._firmware_update_info: return None return self._firmware_update_info.update_available async def get_update_state(self) -> DownloadState: """Return update state.""" resp = await self.call("get_fw_download_state") state = resp["get_fw_download_state"] return DownloadState.from_dict(state) @allow_update_after async def update( self, progress_cb: Callable[[DownloadState], Coroutine] | None = None ) -> dict: """Update the device firmware.""" if not self._firmware_update_info: raise KasaException( "You must call check_latest_firmware before calling update" ) if not self.update_available: raise KasaException("A new update must be available to call update") current_fw = self.current_firmware _LOGGER.info( "Going to upgrade from %s to %s", current_fw, self._firmware_update_info.version, ) await self.call("fw_download") # TODO: read timeout from get_auto_update_info or from get_fw_download_state? async with asyncio_timeout(60 * 5): while True: await asyncio.sleep(0.5) try: state = await self.get_update_state() except Exception as ex: _LOGGER.warning( "Got exception, maybe the device is rebooting? %s", ex ) continue _LOGGER.debug("Update state: %s", state) if progress_cb is not None: asyncio.create_task(progress_cb(state)) if state.status == 0: _LOGGER.info( "Update idle, hopefully updated to %s", self._firmware_update_info.version, ) break elif state.status == 2: _LOGGER.info("Downloading firmware, progress: %s", state.progress) elif state.status == 3: upgrade_sleep = state.upgrade_time _LOGGER.info( "Flashing firmware, sleeping for %s before checking status", upgrade_sleep, ) await asyncio.sleep(upgrade_sleep) elif state.status < 0: _LOGGER.error("Got error: %s", state.status) break else: _LOGGER.warning("Unhandled state code: %s", state) return state.to_dict() @property def auto_update_enabled(self) -> bool: """Return True if autoupdate is enabled.""" return "enable" in self.data and self.data["enable"] @allow_update_after async def set_auto_update_enabled(self, enabled: bool) -> dict: """Change autoupdate setting.""" data = {**self.data, "enable": enabled} return await self.call("set_auto_update_info", data)