Improve encrypted image support. (#153)

* Improve encrypted image support

* Config flow logic for separate encryption key

* Migrate config version.

* Make test RTSP creds optional (Allows for encryption to be used for cameras without RTSP)

custom_components/ezviz_cloud/__init__.py

@@ -12,13 +12,21 @@
 )
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_TIMEOUT, CONF_TYPE, CONF_URL, Platform
+from homeassistant.const import (
+    CONF_PASSWORD,
+    CONF_TIMEOUT,
+    CONF_TYPE,
+    CONF_URL,
+    CONF_USERNAME,
+    Platform,
+)
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 
 from .const import (
     ATTR_TYPE_CAMERA,
     ATTR_TYPE_CLOUD,
+    CONF_ENC_KEY,
     CONF_FFMPEG_ARGUMENTS,
     CONF_RF_SESSION_ID,
     CONF_SESSION_ID,
@@ -67,10 +75,6 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     # Initialize EZVIZ cloud entities
     if PLATFORMS_BY_TYPE[sensor_type]:
-        # Initiate reauth config flow if account token if not present.
-        if not entry.data.get(CONF_SESSION_ID):
-            raise ConfigEntryAuthFailed
-
         ezviz_client = EzvizClient(
             token={
                 CONF_SESSION_ID: entry.data[CONF_SESSION_ID],
@@ -134,3 +138,33 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 async def _async_update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None:
     """Handle options update."""
     await hass.config_entries.async_reload(entry.entry_id)
+
+
+async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+    """Migrate old camera entry."""
+    _LOGGER.debug("Migrating from version %s.%s", entry.version, entry.minor_version)
+
+    if entry.version == 1:
+        if entry.data[CONF_TYPE] == ATTR_TYPE_CAMERA:
+            data = {
+                CONF_USERNAME: entry.data[CONF_USERNAME],
+                CONF_PASSWORD: entry.data[CONF_PASSWORD],
+                CONF_ENC_KEY: entry.data[CONF_PASSWORD],
+                CONF_TYPE: ATTR_TYPE_CAMERA,
+            }
+
+            hass.config_entries.async_update_entry(entry, data=data, version=2)
+
+        if entry.data[CONF_TYPE] == ATTR_TYPE_CLOUD:
+            if not entry.data.get(CONF_SESSION_ID):
+                raise ConfigEntryAuthFailed
+            hass.config_entries.async_update_entry(entry, data=entry.data, version=2)
+
+        _LOGGER.info(
+            "Migration to version %s.%s successful for %s account",
+            entry.version,
+            entry.minor_version,
+            entry.data[CONF_TYPE],
+        )
+
+    return True

custom_components/ezviz_cloud/config_flow.py

@@ -38,9 +38,11 @@
     ATTR_SERIAL,
     ATTR_TYPE_CAMERA,
     ATTR_TYPE_CLOUD,
+    CONF_ENC_KEY,
     CONF_FFMPEG_ARGUMENTS,
     CONF_RF_SESSION_ID,
     CONF_SESSION_ID,
+    CONF_TEST_RTSP_CREDENTIALS,
     DEFAULT_CAMERA_USERNAME,
     DEFAULT_FFMPEG_ARGUMENTS,
     DEFAULT_TIMEOUT,
@@ -84,7 +86,7 @@ def _get_cam_enc_key(data: dict, ezviz_client: EzvizClient) -> Any:
 class EzvizConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for EZVIZ."""
 
-    VERSION = 1
+    VERSION = 2
 
     ip_address: str
     username: str | None
@@ -145,20 +147,25 @@ async def _validate_and_create_camera_rtsp(self, data: dict) -> ConfigFlowResult
         # Create Ezviz API Client.
         await self.hass.async_add_executor_job(ezviz_client.login)
 
-        # If no encryption key is provided, get it. Sometimes a old key is provided and doesn't work.
+        # Fetch encryption key from ezviz api.
+        data[CONF_ENC_KEY] = await self.hass.async_add_executor_job(
+            _get_cam_enc_key, data, ezviz_client
+        )
+
+        # If newer camera, the encryption key is the password.
         if data[CONF_PASSWORD] == "fetch_my_key":
-            data[CONF_PASSWORD] = await self.hass.async_add_executor_job(
-                _get_cam_enc_key, data, ezviz_client
-            )
+            data[CONF_PASSWORD] = data[CONF_ENC_KEY]
 
-        # Test camera RTSP credentials.
-        await self.hass.async_add_executor_job(_wake_camera, data, ezviz_client)
+        # Test camera RTSP credentials. Older cameras still use the verification code on the camera and not the encryption key.
+        if data[CONF_TEST_RTSP_CREDENTIALS]:
+            await self.hass.async_add_executor_job(_wake_camera, data, ezviz_client)
 
         return self.async_create_entry(
             title=data[ATTR_SERIAL],
             data={
                 CONF_USERNAME: data[CONF_USERNAME],
                 CONF_PASSWORD: data[CONF_PASSWORD],
+                CONF_ENC_KEY: data[CONF_ENC_KEY],
                 CONF_TYPE: ATTR_TYPE_CAMERA,
             },
             options=DEFAULT_OPTIONS,
@@ -337,6 +344,7 @@ async def async_step_confirm(
             {
                 vol.Required(CONF_USERNAME, default=DEFAULT_CAMERA_USERNAME): str,
                 vol.Required(CONF_PASSWORD, default="fetch_my_key"): str,
+                vol.Optional(CONF_TEST_RTSP_CREDENTIALS, default=True): bool,
             }
         )
 

custom_components/ezviz_cloud/const.py

@@ -11,6 +11,8 @@
 CONF_SESSION_ID = "session_id"
 CONF_RF_SESSION_ID = "rf_session_id"
 CONF_EZVIZ_ACCOUNT = "ezviz_account"
+CONF_ENC_KEY = "enc_key"
+CONF_TEST_RTSP_CREDENTIALS = "test_rtsp_credentials"
 
 # Service names
 SERVICE_WAKE_DEVICE = "wake_device"

custom_components/ezviz_cloud/image.py

@@ -8,13 +8,12 @@
 from pyezvizapi.utils import decrypt_image
 
 from homeassistant.components.image import Image, ImageEntity, ImageEntityDescription
-from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_PASSWORD
+from homeassistant.config_entries import SOURCE_IGNORE, ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import dt as dt_util
 
-from .const import DATA_COORDINATOR, DOMAIN
+from .const import CONF_ENC_KEY, DATA_COORDINATOR, DOMAIN
 from .coordinator import EzvizDataUpdateCoordinator
 from .entity import EzvizEntity
 
@@ -57,7 +56,9 @@ def __init__(
         )
         camera = hass.config_entries.async_entry_for_domain_unique_id(DOMAIN, serial)
         self.alarm_image_password = (
-            camera.data.get(CONF_PASSWORD) if camera is not None else None
+            camera.data[CONF_ENC_KEY]
+            if camera and camera.source != SOURCE_IGNORE
+            else None
         )
 
     async def _async_load_image_from_url(self, url: str) -> Image | None:

custom_components/ezviz_cloud/manifest.json

@@ -9,5 +9,5 @@
   "issue_tracker": "https://github.com/RenierM26/ha-ezviz/issues",
   "loggers": ["paho_mqtt", "pyezvizapi"],
   "requirements": ["pyezvizapi==1.0.0.0"],
-  "version": "0.1.0.45"
+  "version": "0.1.0.46"
 }

custom_components/ezviz_cloud/strings.json

@@ -24,7 +24,8 @@
         "description": "Enter RTSP credentials for EZVIZ camera {serial} with IP {ip_address}",
         "data": {
           "username": "[%key:common::config_flow::data::username%]",
-          "password": "[%key:common::config_flow::data::password%]"
+          "password": "[%key:common::config_flow::data::password%]",
+          "test_rtsp_credentials": "Test RTSP credentials"
         }
       },
       "reauth_confirm": {