237 lines
9.0 KiB
Python
237 lines
9.0 KiB
Python
"""Representation of ZHA updates."""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import functools
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from zigpy.ota.image import BaseOTAImage
|
|
from zigpy.types import uint16_t
|
|
from zigpy.zcl.foundation import Status
|
|
|
|
from homeassistant.components.update import (
|
|
ATTR_INSTALLED_VERSION,
|
|
ATTR_LATEST_VERSION,
|
|
UpdateDeviceClass,
|
|
UpdateEntity,
|
|
UpdateEntityFeature,
|
|
)
|
|
from homeassistant.config_entries import ConfigEntry
|
|
from homeassistant.const import EntityCategory, Platform
|
|
from homeassistant.core import HomeAssistant, callback
|
|
from homeassistant.exceptions import HomeAssistantError
|
|
from homeassistant.helpers.dispatcher import async_dispatcher_connect
|
|
from homeassistant.helpers.entity_platform import AddEntitiesCallback
|
|
from homeassistant.helpers.restore_state import ExtraStoredData
|
|
|
|
from .core import discovery
|
|
from .core.const import CLUSTER_HANDLER_OTA, SIGNAL_ADD_ENTITIES, UNKNOWN
|
|
from .core.helpers import get_zha_data
|
|
from .core.registries import ZHA_ENTITIES
|
|
from .entity import ZhaEntity
|
|
|
|
if TYPE_CHECKING:
|
|
from .core.cluster_handlers import ClusterHandler
|
|
from .core.device import ZHADevice
|
|
|
|
CONFIG_DIAGNOSTIC_MATCH = functools.partial(
|
|
ZHA_ENTITIES.config_diagnostic_match, Platform.UPDATE
|
|
)
|
|
|
|
# don't let homeassistant check for updates button hammer the zigbee network
|
|
PARALLEL_UPDATES = 1
|
|
|
|
|
|
@dataclass
|
|
class ZHAFirmwareUpdateExtraStoredData(ExtraStoredData):
|
|
"""Extra stored data for ZHA firmware update entity."""
|
|
|
|
image_type: uint16_t | None
|
|
|
|
def as_dict(self) -> dict[str, Any]:
|
|
"""Return a dict representation of the extra data."""
|
|
return {"image_type": self.image_type}
|
|
|
|
|
|
async def async_setup_entry(
|
|
hass: HomeAssistant,
|
|
config_entry: ConfigEntry,
|
|
async_add_entities: AddEntitiesCallback,
|
|
) -> None:
|
|
"""Set up the Zigbee Home Automation update from config entry."""
|
|
zha_data = get_zha_data(hass)
|
|
entities_to_create = zha_data.platforms[Platform.UPDATE]
|
|
|
|
unsub = async_dispatcher_connect(
|
|
hass,
|
|
SIGNAL_ADD_ENTITIES,
|
|
functools.partial(
|
|
discovery.async_add_entities, async_add_entities, entities_to_create
|
|
),
|
|
)
|
|
config_entry.async_on_unload(unsub)
|
|
|
|
|
|
@CONFIG_DIAGNOSTIC_MATCH(cluster_handler_names=CLUSTER_HANDLER_OTA)
|
|
class ZHAFirmwareUpdateEntity(ZhaEntity, UpdateEntity):
|
|
"""Representation of a ZHA firmware update entity."""
|
|
|
|
_unique_id_suffix = "firmware_update"
|
|
_attr_entity_category = EntityCategory.CONFIG
|
|
_attr_device_class = UpdateDeviceClass.FIRMWARE
|
|
_attr_supported_features = (
|
|
UpdateEntityFeature.INSTALL
|
|
| UpdateEntityFeature.PROGRESS
|
|
| UpdateEntityFeature.SPECIFIC_VERSION
|
|
)
|
|
|
|
def __init__(
|
|
self,
|
|
unique_id: str,
|
|
zha_device: ZHADevice,
|
|
channels: list[ClusterHandler],
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""Initialize the ZHA update entity."""
|
|
super().__init__(unique_id, zha_device, channels, **kwargs)
|
|
self._ota_cluster_handler: ClusterHandler = self.cluster_handlers[
|
|
CLUSTER_HANDLER_OTA
|
|
]
|
|
self._attr_installed_version: str = self.determine_installed_version()
|
|
self._image_type: uint16_t | None = None
|
|
self._latest_version_firmware: BaseOTAImage | None = None
|
|
self._result = None
|
|
|
|
@callback
|
|
def determine_installed_version(self) -> str:
|
|
"""Determine the currently installed firmware version."""
|
|
currently_installed_version = self._ota_cluster_handler.current_file_version
|
|
version_from_dr = self.zha_device.sw_version
|
|
if currently_installed_version == UNKNOWN and version_from_dr:
|
|
currently_installed_version = version_from_dr
|
|
return currently_installed_version
|
|
|
|
@property
|
|
def extra_restore_state_data(self) -> ZHAFirmwareUpdateExtraStoredData:
|
|
"""Return ZHA firmware update specific state data to be restored."""
|
|
return ZHAFirmwareUpdateExtraStoredData(self._image_type)
|
|
|
|
@callback
|
|
def device_ota_update_available(self, image: BaseOTAImage) -> None:
|
|
"""Handle ota update available signal from Zigpy."""
|
|
self._latest_version_firmware = image
|
|
self._attr_latest_version = f"0x{image.header.file_version:08x}"
|
|
self._image_type = image.header.image_type
|
|
self._attr_installed_version = self.determine_installed_version()
|
|
self.async_write_ha_state()
|
|
|
|
@callback
|
|
def _update_progress(self, current: int, total: int, progress: float) -> None:
|
|
"""Update install progress on event."""
|
|
assert self._latest_version_firmware
|
|
self._attr_in_progress = int(progress)
|
|
self.async_write_ha_state()
|
|
|
|
@callback
|
|
def _reset_progress(self, write_state: bool = True) -> None:
|
|
"""Reset update install progress."""
|
|
self._result = None
|
|
self._attr_in_progress = False
|
|
if write_state:
|
|
self.async_write_ha_state()
|
|
|
|
async def async_update(self) -> None:
|
|
"""Handle the update entity service call to manually check for available firmware updates."""
|
|
await super().async_update()
|
|
# check for updates in the HA settings menu can invoke this so we need to check if the device
|
|
# is mains powered so we don't get a ton of errors in the logs from sleepy devices.
|
|
if self.zha_device.available and self.zha_device.is_mains_powered:
|
|
await self._ota_cluster_handler.async_check_for_update()
|
|
|
|
async def async_install(
|
|
self, version: str | None, backup: bool, **kwargs: Any
|
|
) -> None:
|
|
"""Install an update."""
|
|
firmware = self._latest_version_firmware
|
|
assert firmware
|
|
self._reset_progress(False)
|
|
self._attr_in_progress = True
|
|
self.async_write_ha_state()
|
|
|
|
try:
|
|
self._result = await self.zha_device.device.update_firmware(
|
|
self._latest_version_firmware,
|
|
self._update_progress,
|
|
)
|
|
except Exception as ex:
|
|
self._reset_progress()
|
|
raise HomeAssistantError(ex) from ex
|
|
|
|
assert self._result is not None
|
|
|
|
# If the update was not successful, we should throw an error to let the user know
|
|
if self._result != Status.SUCCESS:
|
|
# save result since reset_progress will clear it
|
|
results = self._result
|
|
self._reset_progress()
|
|
raise HomeAssistantError(f"Update was not successful - result: {results}")
|
|
|
|
# If we get here, all files were installed successfully
|
|
self._attr_installed_version = (
|
|
self._attr_latest_version
|
|
) = f"0x{firmware.header.file_version:08x}"
|
|
self._latest_version_firmware = None
|
|
self._reset_progress()
|
|
|
|
async def async_added_to_hass(self) -> None:
|
|
"""Call when entity is added."""
|
|
await super().async_added_to_hass()
|
|
last_state = await self.async_get_last_state()
|
|
# If we have a complete previous state, use that to set the installed version
|
|
if (
|
|
last_state
|
|
and self._attr_installed_version == UNKNOWN
|
|
and (installed_version := last_state.attributes.get(ATTR_INSTALLED_VERSION))
|
|
):
|
|
self._attr_installed_version = installed_version
|
|
# If we have a complete previous state, use that to set the latest version
|
|
if (
|
|
last_state
|
|
and (latest_version := last_state.attributes.get(ATTR_LATEST_VERSION))
|
|
is not None
|
|
and latest_version != UNKNOWN
|
|
):
|
|
self._attr_latest_version = latest_version
|
|
# If we have no state or latest version to restore, or the latest version is
|
|
# the same as the installed version, we can set the latest
|
|
# version to installed so that the entity starts as off.
|
|
elif (
|
|
not last_state
|
|
or not latest_version
|
|
or latest_version == self._attr_installed_version
|
|
):
|
|
self._attr_latest_version = self._attr_installed_version
|
|
|
|
if self._attr_latest_version != self._attr_installed_version and (
|
|
extra_data := await self.async_get_last_extra_data()
|
|
):
|
|
self._image_type = extra_data.as_dict()["image_type"]
|
|
if self._image_type:
|
|
self._latest_version_firmware = (
|
|
await self.zha_device.device.application.ota.get_ota_image(
|
|
self.zha_device.manufacturer_code, self._image_type
|
|
)
|
|
)
|
|
# if we can't locate an image but we have a latest version that differs
|
|
# we should set the latest version to the installed version to avoid
|
|
# confusion and errors
|
|
if not self._latest_version_firmware:
|
|
self._attr_latest_version = self._attr_installed_version
|
|
|
|
self.zha_device.device.add_listener(self)
|
|
|
|
async def async_will_remove_from_hass(self) -> None:
|
|
"""Call when entity will be removed."""
|
|
await super().async_will_remove_from_hass()
|
|
self._reset_progress(False)
|