Add option to retrieve SwitchBot Lock encryption key through config flow (#84830)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/switchbot/config_flow.py

@@ -17,7 +17,12 @@ from homeassistant.components.bluetooth import (
     async_discovered_service_info,
 )
 from homeassistant.config_entries import ConfigEntry, ConfigFlow, OptionsFlow
-from homeassistant.const import CONF_ADDRESS, CONF_PASSWORD, CONF_SENSOR_TYPE
+from homeassistant.const import (
+    CONF_ADDRESS,
+    CONF_PASSWORD,
+    CONF_SENSOR_TYPE,
+    CONF_USERNAME,
+)
 from homeassistant.core import callback
 from homeassistant.data_entry_flow import AbortFlow, FlowResult
 
@@ -94,6 +99,8 @@ class SwitchbotConfigFlow(ConfigFlow, domain=DOMAIN):
             "name": data["modelFriendlyName"],
             "address": short_address(discovery_info.address),
         }
+        if model_name == SwitchbotModel.LOCK:
+            return await self.async_step_lock_chose_method()
         if self._discovered_adv.data["isEncrypted"]:
             return await self.async_step_password()
         return await self.async_step_confirm()
@@ -151,6 +158,57 @@ class SwitchbotConfigFlow(ConfigFlow, domain=DOMAIN):
             },
         )
 
+    async def async_step_lock_auth(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle the SwitchBot API auth step."""
+        errors = {}
+        assert self._discovered_adv is not None
+        if user_input is not None:
+            try:
+                key_details = await self.hass.async_add_executor_job(
+                    SwitchbotLock.retrieve_encryption_key,
+                    self._discovered_adv.address,
+                    user_input[CONF_USERNAME],
+                    user_input[CONF_PASSWORD],
+                )
+                return await self.async_step_lock_key(key_details)
+            except RuntimeError:
+                errors = {
+                    "base": "auth_failed",
+                }
+
+        user_input = user_input or {}
+        return self.async_show_form(
+            step_id="lock_auth",
+            errors=errors,
+            data_schema=vol.Schema(
+                {
+                    vol.Required(
+                        CONF_USERNAME, default=user_input.get(CONF_USERNAME)
+                    ): str,
+                    vol.Required(CONF_PASSWORD): str,
+                }
+            ),
+            description_placeholders={
+                "name": name_from_discovery(self._discovered_adv),
+            },
+        )
+
+    async def async_step_lock_chose_method(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle the SwitchBot API chose method step."""
+        assert self._discovered_adv is not None
+
+        return self.async_show_menu(
+            step_id="lock_chose_method",
+            menu_options=["lock_auth", "lock_key"],
+            description_placeholders={
+                "name": name_from_discovery(self._discovered_adv),
+            },
+        )
+
     async def async_step_lock_key(
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
@@ -160,12 +218,11 @@ class SwitchbotConfigFlow(ConfigFlow, domain=DOMAIN):
         if user_input is not None:
             if not await SwitchbotLock.verify_encryption_key(
                 self._discovered_adv.device,
-                user_input.get(CONF_KEY_ID),
+                user_input[CONF_KEY_ID],
-                user_input.get(CONF_ENCRYPTION_KEY),
+                user_input[CONF_ENCRYPTION_KEY],
             ):
                 errors = {
-                    CONF_KEY_ID: "key_id_invalid",
+                    "base": "encryption_key_invalid",
-                    CONF_ENCRYPTION_KEY: "encryption_key_invalid",
                 }
             else:
                 return await self._async_create_entry_from_discovery(user_input)
@@ -229,7 +286,7 @@ class SwitchbotConfigFlow(ConfigFlow, domain=DOMAIN):
             device_adv = self._discovered_advs[user_input[CONF_ADDRESS]]
             await self._async_set_device(device_adv)
             if device_adv.data.get("modelName") == SwitchbotModel.LOCK:
-                return await self.async_step_lock_key()
+                return await self.async_step_lock_chose_method()
             if device_adv.data["isEncrypted"]:
                 return await self.async_step_password()
             return await self._async_create_entry_from_discovery(user_input)
@@ -241,7 +298,7 @@ class SwitchbotConfigFlow(ConfigFlow, domain=DOMAIN):
             device_adv = list(self._discovered_advs.values())[0]
             await self._async_set_device(device_adv)
             if device_adv.data.get("modelName") == SwitchbotModel.LOCK:
-                return await self.async_step_lock_key()
+                return await self.async_step_lock_chose_method()
             if device_adv.data["isEncrypted"]:
                 return await self.async_step_password()
             return await self.async_step_confirm()

homeassistant/components/switchbot/manifest.json

@@ -2,7 +2,7 @@
   "domain": "switchbot",
   "name": "SwitchBot",
   "documentation": "https://www.home-assistant.io/integrations/switchbot",
-  "requirements": ["PySwitchbot==0.33.0"],
+  "requirements": ["PySwitchbot==0.34.1"],
   "config_flow": true,
   "dependencies": ["bluetooth"],
   "codeowners": [

homeassistant/components/switchbot/strings.json

@@ -22,11 +22,25 @@
           "key_id": "Key ID",
           "encryption_key": "Encryption key"
         }
+      },
+      "lock_auth": {
+        "description": "Please provide your SwitchBot app username and password. This data won't be saved and only used to retrieve your locks encryption key.",
+        "data": {
+          "username": "[%key:common::config_flow::data::username%]",
+          "password": "[%key:common::config_flow::data::password%]"
+        }
+      },
+      "lock_chose_method": {
+        "description": "Choose configuration method, details can be found in the documentation.",
+        "menu_options": {
+          "lock_auth": "SwitchBot app login and password",
+          "lock_key": "Lock encryption key"
+        }
       }
     },
     "error": {
-      "key_id_invalid": "Key ID or Encryption key is invalid",
+      "encryption_key_invalid": "Key ID or Encryption key is invalid",
-      "encryption_key_invalid": "Key ID or Encryption key is invalid"
+      "auth_failed": "Authentication failed"
     },
     "abort": {
       "already_configured_device": "[%key:common::config_flow::abort::already_configured_device%]",

homeassistant/components/switchbot/translations/en.json

@@ -7,11 +7,36 @@
             "switchbot_unsupported_type": "Unsupported Switchbot Type.",
             "unknown": "Unexpected error"
         },
+        "error": {
+            "auth_failed": "Authentication failed",
+            "encryption_key_invalid": "Key ID or Encryption key is invalid"
+        },
         "flow_title": "{name} ({address})",
         "step": {
             "confirm": {
                 "description": "Do you want to set up {name}?"
             },
+            "lock_auth": {
+                "data": {
+                    "password": "Password",
+                    "username": "Username"
+                },
+                "description": "Please provide your SwitchBot app username and password. This data won't be saved and only used to retrieve your locks encryption key."
+            },
+            "lock_chose_method": {
+                "description": "Choose configuration method, details can be found in the documentation.",
+                "menu_options": {
+                    "lock_auth": "SwitchBot app login and password",
+                    "lock_key": "Lock encryption key"
+                }
+            },
+            "lock_key": {
+                "data": {
+                    "encryption_key": "Encryption key",
+                    "key_id": "Key ID"
+                },
+                "description": "The {name} device requires encryption key, details on how to obtain it can be found in the documentation."
+            },
             "password": {
                 "data": {
                     "password": "Password"

requirements_all.txt

@@ -40,7 +40,7 @@ PyRMVtransport==0.3.3
 PySocks==1.7.1
 
 # homeassistant.components.switchbot
-PySwitchbot==0.33.0
+PySwitchbot==0.34.1
 
 # homeassistant.components.transport_nsw
 PyTransportNSW==0.1.1

requirements_test_all.txt

@@ -36,7 +36,7 @@ PyRMVtransport==0.3.3
 PySocks==1.7.1
 
 # homeassistant.components.switchbot
-PySwitchbot==0.33.0
+PySwitchbot==0.34.1
 
 # homeassistant.components.transport_nsw
 PyTransportNSW==0.1.1

tests/components/switchbot/__init__.py

@@ -180,7 +180,7 @@ WOLOCK_SERVICE_INFO = BluetoothServiceInfoBleak(
     advertisement=generate_advertisement_data(
         local_name="WoLock",
         manufacturer_data={2409: b"\xf1\t\x9fE\x1a]\xda\x83\x00 "},
-        service_data={"00000d00-0000-1000-8000-00805f9b34fb": b"o\x80d"},
+        service_data={"0000fd3d-0000-1000-8000-00805f9b34fb": b"o\x80d"},
         service_uuids=["cba20d00-224d-11e6-9fb8-0002a5d5c51b"],
     ),
     device=BLEDevice("aa:bb:cc:dd:ee:ff", "WoLock"),

tests/components/switchbot/test_config_flow.py

@@ -8,7 +8,13 @@ from homeassistant.components.switchbot.const import (
     CONF_RETRY_COUNT,
 )
 from homeassistant.config_entries import SOURCE_BLUETOOTH, SOURCE_USER
-from homeassistant.const import CONF_ADDRESS, CONF_NAME, CONF_PASSWORD, CONF_SENSOR_TYPE
+from homeassistant.const import (
+    CONF_ADDRESS,
+    CONF_NAME,
+    CONF_PASSWORD,
+    CONF_SENSOR_TYPE,
+    CONF_USERNAME,
+)
 from homeassistant.data_entry_flow import FlowResultType
 
 from . import (
@@ -85,6 +91,66 @@ async def test_bluetooth_discovery_requires_password(hass):
     assert len(mock_setup_entry.mock_calls) == 1
 
 
+async def test_bluetooth_discovery_lock_key(hass):
+    """Test discovery via bluetooth with a lock."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        context={"source": SOURCE_BLUETOOTH},
+        data=WOLOCK_SERVICE_INFO,
+    )
+    assert result["type"] == FlowResultType.MENU
+    assert result["step_id"] == "lock_chose_method"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], user_input={"next_step_id": "lock_key"}
+    )
+    await hass.async_block_till_done()
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "lock_key"
+    assert result["errors"] == {}
+
+    with patch(
+        "homeassistant.components.switchbot.config_flow.SwitchbotLock.verify_encryption_key",
+        return_value=False,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                CONF_KEY_ID: "",
+                CONF_ENCRYPTION_KEY: "",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "lock_key"
+    assert result["errors"] == {"base": "encryption_key_invalid"}
+
+    with patch_async_setup_entry() as mock_setup_entry, patch(
+        "homeassistant.components.switchbot.config_flow.SwitchbotLock.verify_encryption_key",
+        return_value=True,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                CONF_KEY_ID: "ff",
+                CONF_ENCRYPTION_KEY: "ffffffffffffffffffffffffffffffff",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == FlowResultType.CREATE_ENTRY
+    assert result["title"] == "Lock EEFF"
+    assert result["data"] == {
+        CONF_ADDRESS: "aa:bb:cc:dd:ee:ff",
+        CONF_KEY_ID: "ff",
+        CONF_ENCRYPTION_KEY: "ffffffffffffffffffffffffffffffff",
+        CONF_SENSOR_TYPE: "lock",
+    }
+
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
 async def test_bluetooth_discovery_already_setup(hass):
     """Test discovery via bluetooth with a valid device when already setup."""
     entry = MockConfigEntry(
@@ -327,7 +393,7 @@ async def test_user_setup_single_bot_with_password(hass):
     assert len(mock_setup_entry.mock_calls) == 1
 
 
-async def test_user_setup_wolock(hass):
+async def test_user_setup_wolock_key(hass):
     """Test the user initiated form for a lock."""
 
     with patch(
@@ -337,14 +403,39 @@ async def test_user_setup_wolock(hass):
         result = await hass.config_entries.flow.async_init(
             DOMAIN, context={"source": SOURCE_USER}
         )
+    assert result["type"] == FlowResultType.MENU
+    assert result["step_id"] == "lock_chose_method"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], user_input={"next_step_id": "lock_key"}
+    )
+    await hass.async_block_till_done()
     assert result["type"] == FlowResultType.FORM
     assert result["step_id"] == "lock_key"
     assert result["errors"] == {}
 
-    with patch_async_setup_entry() as mock_setup_entry, patch(
+    with patch(
-        "switchbot.SwitchbotLock.verify_encryption_key", return_value=True
+        "homeassistant.components.switchbot.config_flow.SwitchbotLock.verify_encryption_key",
+        return_value=False,
     ):
-        result2 = await hass.config_entries.flow.async_configure(
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                CONF_KEY_ID: "",
+                CONF_ENCRYPTION_KEY: "",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "lock_key"
+    assert result["errors"] == {"base": "encryption_key_invalid"}
+
+    with patch_async_setup_entry() as mock_setup_entry, patch(
+        "homeassistant.components.switchbot.config_flow.SwitchbotLock.verify_encryption_key",
+        return_value=True,
+    ):
+        result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
             {
                 CONF_KEY_ID: "ff",
@@ -353,9 +444,77 @@ async def test_user_setup_wolock(hass):
         )
         await hass.async_block_till_done()
 
-    assert result2["type"] == FlowResultType.CREATE_ENTRY
+    assert result["type"] == FlowResultType.CREATE_ENTRY
-    assert result2["title"] == "Lock EEFF"
+    assert result["title"] == "Lock EEFF"
-    assert result2["data"] == {
+    assert result["data"] == {
+        CONF_ADDRESS: "aa:bb:cc:dd:ee:ff",
+        CONF_KEY_ID: "ff",
+        CONF_ENCRYPTION_KEY: "ffffffffffffffffffffffffffffffff",
+        CONF_SENSOR_TYPE: "lock",
+    }
+
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_user_setup_wolock_auth(hass):
+    """Test the user initiated form for a lock."""
+
+    with patch(
+        "homeassistant.components.switchbot.config_flow.async_discovered_service_info",
+        return_value=[WOLOCK_SERVICE_INFO],
+    ):
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": SOURCE_USER}
+        )
+    assert result["type"] == FlowResultType.MENU
+    assert result["step_id"] == "lock_chose_method"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], user_input={"next_step_id": "lock_auth"}
+    )
+    await hass.async_block_till_done()
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "lock_auth"
+    assert result["errors"] == {}
+
+    with patch(
+        "homeassistant.components.switchbot.config_flow.SwitchbotLock.retrieve_encryption_key",
+        side_effect=RuntimeError,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                CONF_USERNAME: "",
+                CONF_PASSWORD: "",
+            },
+        )
+        await hass.async_block_till_done()
+    assert result["type"] == FlowResultType.FORM
+    assert result["step_id"] == "lock_auth"
+    assert result["errors"] == {"base": "auth_failed"}
+
+    with patch_async_setup_entry() as mock_setup_entry, patch(
+        "homeassistant.components.switchbot.config_flow.SwitchbotLock.verify_encryption_key",
+        return_value=True,
+    ), patch(
+        "homeassistant.components.switchbot.config_flow.SwitchbotLock.retrieve_encryption_key",
+        return_value={
+            CONF_KEY_ID: "ff",
+            CONF_ENCRYPTION_KEY: "ffffffffffffffffffffffffffffffff",
+        },
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                CONF_USERNAME: "username",
+                CONF_PASSWORD: "password",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == FlowResultType.CREATE_ENTRY
+    assert result["title"] == "Lock EEFF"
+    assert result["data"] == {
         CONF_ADDRESS: "aa:bb:cc:dd:ee:ff",
         CONF_KEY_ID: "ff",
         CONF_ENCRYPTION_KEY: "ffffffffffffffffffffffffffffffff",
@@ -387,12 +546,20 @@ async def test_user_setup_wolock_or_bot(hass):
         USER_INPUT,
     )
     await hass.async_block_till_done()
+    assert result["type"] == FlowResultType.MENU
+    assert result["step_id"] == "lock_chose_method"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], user_input={"next_step_id": "lock_key"}
+    )
+    await hass.async_block_till_done()
     assert result["type"] == FlowResultType.FORM
     assert result["step_id"] == "lock_key"
     assert result["errors"] == {}
 
     with patch_async_setup_entry() as mock_setup_entry, patch(
-        "switchbot.SwitchbotLock.verify_encryption_key", return_value=True
+        "homeassistant.components.switchbot.config_flow.SwitchbotLock.verify_encryption_key",
+        return_value=True,
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
@@ -415,42 +582,6 @@ async def test_user_setup_wolock_or_bot(hass):
     assert len(mock_setup_entry.mock_calls) == 1
 
 
-async def test_user_setup_wolock_invalid_encryption_key(hass):
-    """Test the user initiated form for a lock with invalid encryption key."""
-
-    with patch(
-        "homeassistant.components.switchbot.config_flow.async_discovered_service_info",
-        return_value=[WOLOCK_SERVICE_INFO],
-    ):
-        result = await hass.config_entries.flow.async_init(
-            DOMAIN, context={"source": SOURCE_USER}
-        )
-    assert result["type"] == FlowResultType.FORM
-    assert result["step_id"] == "lock_key"
-    assert result["errors"] == {}
-
-    with patch_async_setup_entry() as mock_setup_entry, patch(
-        "switchbot.SwitchbotLock.verify_encryption_key", return_value=False
-    ):
-        result2 = await hass.config_entries.flow.async_configure(
-            result["flow_id"],
-            {
-                CONF_KEY_ID: "",
-                CONF_ENCRYPTION_KEY: "",
-            },
-        )
-        await hass.async_block_till_done()
-
-    assert result2["type"] == FlowResultType.FORM
-    assert result2["step_id"] == "lock_key"
-    assert result2["errors"] == {
-        CONF_KEY_ID: "key_id_invalid",
-        CONF_ENCRYPTION_KEY: "encryption_key_invalid",
-    }
-
-    assert len(mock_setup_entry.mock_calls) == 0
-
-
 async def test_user_setup_wosensor(hass):
     """Test the user initiated form with password and valid mac."""
     with patch(