Add security options to disable AP and BLE RPC after Shelly WiFi provisioning (#156970)

homeassistant/components/shelly/config_flow.py

@@ -183,6 +183,8 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN):
     selected_ssid: str = ""
     _provision_task: asyncio.Task | None = None
     _provision_result: ConfigFlowResult | None = None
+    disable_ap_after_provision: bool = True
+    disable_ble_rpc_after_provision: bool = True
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
@@ -426,10 +428,20 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN):
     ) -> ConfigFlowResult:
         """Confirm bluetooth provisioning."""
         if user_input is not None:
+            self.disable_ap_after_provision = user_input.get("disable_ap", True)
+            self.disable_ble_rpc_after_provision = user_input.get(
+                "disable_ble_rpc", True
+            )
             return await self.async_step_wifi_scan()
 
         return self.async_show_form(
             step_id="bluetooth_confirm",
+            data_schema=vol.Schema(
+                {
+                    vol.Optional("disable_ap", default=True): bool,
+                    vol.Optional("disable_ble_rpc", default=True): bool,
+                }
+            ),
             description_placeholders={
                 "name": self.context["title_placeholders"]["name"]
             },
@@ -521,6 +533,62 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN):
             description_placeholders={"ssid": self.selected_ssid},
         )
 
+    async def _async_secure_device_after_provision(self, host: str, port: int) -> None:
+        """Disable AP and/or BLE RPC after successful WiFi provisioning.
+
+        Must be called via IP after device is on WiFi, not via BLE.
+        """
+        if (
+            not self.disable_ap_after_provision
+            and not self.disable_ble_rpc_after_provision
+        ):
+            return
+
+        # Connect to device via IP
+        options = ConnectionOptions(
+            host,
+            None,
+            None,
+            device_mac=self.unique_id,
+            port=port,
+        )
+        device: RpcDevice | None = None
+        try:
+            device = await RpcDevice.create(
+                async_get_clientsession(self.hass), None, options
+            )
+            await device.initialize()
+
+            restart_required = False
+
+            # Disable WiFi AP if requested
+            if self.disable_ap_after_provision:
+                result = await device.wifi_setconfig(ap_enable=False)
+                LOGGER.debug("Disabled WiFi AP on %s", host)
+                restart_required = restart_required or result.get(
+                    "restart_required", False
+                )
+
+            # Disable BLE RPC if requested (keep BLE enabled for sensors/buttons)
+            if self.disable_ble_rpc_after_provision:
+                result = await device.ble_setconfig(enable=True, enable_rpc=False)
+                LOGGER.debug("Disabled BLE RPC on %s", host)
+                restart_required = restart_required or result.get(
+                    "restart_required", False
+                )
+
+            # Restart device once if either operation requires it
+            if restart_required:
+                await device.trigger_reboot(delay_ms=1000)
+        except (TimeoutError, DeviceConnectionError, RpcCallError) as err:
+            LOGGER.warning(
+                "Failed to secure device after provisioning at %s: %s", host, err
+            )
+            # Don't fail the flow - device is already on WiFi and functional
+        finally:
+            if device:
+                await device.shutdown()
+
     async def _async_provision_wifi_and_wait_for_zeroconf(
         self, mac: str, password: str, state: ProvisioningState
     ) -> ConfigFlowResult | None:
@@ -614,6 +682,9 @@ class ShellyConfigFlow(ConfigFlow, domain=DOMAIN):
         if not device_info[CONF_MODEL]:
             return self.async_abort(reason="firmware_not_fully_provisioned")
 
+        # Secure device after provisioning if requested (disable AP/BLE)
+        await self._async_secure_device_after_provision(self.host, self.port)
+
         # User just provisioned this device - create entry directly without confirmation
         return self.async_create_entry(
             title=device_info["title"],

homeassistant/components/shelly/strings.json

@@ -31,7 +31,16 @@
     },
     "step": {
       "bluetooth_confirm": {
-        "description": "The Shelly device {name} has been discovered via Bluetooth but is not connected to WiFi.\n\nDo you want to provision WiFi credentials to this device?"
+        "data": {
+          "disable_ap": "Disable WiFi access point after provisioning",
+          "disable_ble_rpc": "Disable Bluetooth RPC after provisioning"
+        },
+        "data_description": {
+          "disable_ap": "For improved security, disable the WiFi access point after successfully connecting to your network.",
+          "disable_ble_rpc": "For improved security, disable Bluetooth RPC access after WiFi is configured. Bluetooth will remain enabled for BLE sensors and buttons."
+        },
+        "description": "The Shelly device {name} has been discovered via Bluetooth but is not connected to WiFi.\n\nDo you want to provision WiFi credentials to this device?",
+        "title": "Provision WiFi via Bluetooth"
       },
       "confirm_discovery": {
         "description": "Do you want to set up the {model} at {host}?\n\nBattery-powered devices that are password-protected must be woken up before continuing with setting up.\nBattery-powered devices that are not password-protected will be added when the device wakes up, you can now manually wake the device up using a button on it or wait for the next data update from the device."

tests/components/shelly/conftest.py

@@ -576,6 +576,9 @@ def _mock_rpc_device(version: str | None = None):
         zigbee_enabled=False,
         zigbee_firmware=False,
         ip_address="10.10.10.10",
+        wifi_setconfig=AsyncMock(return_value={}),
+        ble_setconfig=AsyncMock(return_value={"restart_required": False}),
+        shutdown=AsyncMock(),
     )
     type(device).name = PropertyMock(return_value="Test name")
     return device

tests/components/shelly/test_config_flow.py

@@ -13,6 +13,7 @@ from aioshelly.exceptions import (
     DeviceConnectionError,
     InvalidAuthError,
     InvalidHostError,
+    RpcCallError,
 )
 import pytest
 
@@ -2326,6 +2327,16 @@ async def test_bluetooth_wifi_scan_failure(
     )
 
     # Complete provisioning
+    mock_device = AsyncMock()
+    mock_device.initialize = AsyncMock()
+    mock_device.name = "Test name"
+    mock_device.status = {"sys": {}}
+    mock_device.xmod_info = {}
+    mock_device.shelly = {"model": MODEL_PLUS_2PM}
+    mock_device.wifi_setconfig = AsyncMock(return_value={})
+    mock_device.ble_setconfig = AsyncMock(return_value={"restart_required": False})
+    mock_device.shutdown = AsyncMock()
+
     with (
         patch(
             "homeassistant.components.shelly.config_flow.async_provision_wifi",
@@ -2338,6 +2349,10 @@ async def test_bluetooth_wifi_scan_failure(
             "homeassistant.components.shelly.config_flow.get_info",
             return_value=MOCK_DEVICE_INFO,
         ),
+        patch(
+            "homeassistant.components.shelly.config_flow.RpcDevice.create",
+            return_value=mock_device,
+        ),
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
@@ -2431,6 +2446,16 @@ async def test_bluetooth_wifi_credentials_and_provision_success(
     assert result["step_id"] == "wifi_credentials"
 
     # Enter password and provision
+    mock_device = AsyncMock()
+    mock_device.initialize = AsyncMock()
+    mock_device.name = "Test name"
+    mock_device.status = {"sys": {}}
+    mock_device.xmod_info = {}
+    mock_device.shelly = {"model": MODEL_PLUS_2PM}
+    mock_device.wifi_setconfig = AsyncMock(return_value={})
+    mock_device.ble_setconfig = AsyncMock(return_value={"restart_required": False})
+    mock_device.shutdown = AsyncMock()
+
     with (
         patch(
             "homeassistant.components.shelly.config_flow.async_provision_wifi",
@@ -2443,6 +2468,10 @@ async def test_bluetooth_wifi_credentials_and_provision_success(
             "homeassistant.components.shelly.config_flow.get_info",
             return_value=MOCK_DEVICE_INFO,
         ),
+        patch(
+            "homeassistant.components.shelly.config_flow.RpcDevice.create",
+            return_value=mock_device,
+        ),
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
@@ -2999,6 +3028,17 @@ async def test_bluetooth_provision_with_zeroconf_discovery_fast_path(
         # Ensure the zeroconf discovery completes before returning
         await hass.async_block_till_done()
 
+    # Mock device for secure device feature
+    mock_device = AsyncMock()
+    mock_device.initialize = AsyncMock()
+    mock_device.name = "Test name"
+    mock_device.status = {"sys": {}}
+    mock_device.xmod_info = {}
+    mock_device.shelly = {"model": MODEL_PLUS_2PM}
+    mock_device.wifi_setconfig = AsyncMock(return_value={})
+    mock_device.ble_setconfig = AsyncMock(return_value={"restart_required": False})
+    mock_device.shutdown = AsyncMock()
+
     with (
         patch(
             "homeassistant.components.shelly.config_flow.PROVISIONING_TIMEOUT",
@@ -3016,6 +3056,10 @@ async def test_bluetooth_provision_with_zeroconf_discovery_fast_path(
             "homeassistant.components.shelly.config_flow.get_info",
             return_value=MOCK_DEVICE_INFO,
         ),
+        patch(
+            "homeassistant.components.shelly.config_flow.RpcDevice.create",
+            return_value=mock_device,
+        ),
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],
@@ -3102,6 +3146,390 @@ async def test_bluetooth_provision_timeout_active_lookup_fails(
     assert result["reason"] == "unknown"
 
 
+async def test_bluetooth_provision_secure_device_both_enabled(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_setup: AsyncMock,
+) -> None:
+    """Test provisioning with both AP and BLE disable enabled (default)."""
+    inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        data=BLE_DISCOVERY_INFO,
+        context={"source": config_entries.SOURCE_BLUETOOTH},
+    )
+
+    # Confirm with both switches enabled (default)
+    with patch(
+        "homeassistant.components.shelly.config_flow.async_scan_wifi_networks",
+        return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}],
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"disable_ap": True, "disable_ble_rpc": True},
+        )
+
+    # Select network
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {CONF_SSID: "MyNetwork"},
+    )
+
+    # Provision and verify security calls
+    mock_device = AsyncMock()
+    mock_device.initialize = AsyncMock()
+    mock_device.wifi_setconfig = AsyncMock(return_value={})
+    mock_device.ble_setconfig = AsyncMock(return_value={"restart_required": False})
+    mock_device.shutdown = AsyncMock()
+
+    with (
+        patch("homeassistant.components.shelly.config_flow.async_provision_wifi"),
+        patch(
+            "homeassistant.components.shelly.config_flow.async_lookup_device_by_name",
+            return_value=("1.1.1.1", 80),
+        ),
+        patch(
+            "homeassistant.components.shelly.config_flow.get_info",
+            return_value=MOCK_DEVICE_INFO,
+        ),
+        patch(
+            "homeassistant.components.shelly.config_flow.RpcDevice.create",
+            return_value=mock_device,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_PASSWORD: "my_password"},
+        )
+        await hass.async_block_till_done()
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    # Verify entry created
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+
+    # Verify security calls were made
+    mock_device.wifi_setconfig.assert_called_once_with(ap_enable=False)
+    mock_device.ble_setconfig.assert_called_once_with(enable=True, enable_rpc=False)
+    assert mock_device.shutdown.called
+
+
+async def test_bluetooth_provision_secure_device_both_disabled(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_setup: AsyncMock,
+) -> None:
+    """Test provisioning with both AP and BLE disable disabled."""
+    inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        data=BLE_DISCOVERY_INFO,
+        context={"source": config_entries.SOURCE_BLUETOOTH},
+    )
+
+    # Confirm with both switches disabled
+    with patch(
+        "homeassistant.components.shelly.config_flow.async_scan_wifi_networks",
+        return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}],
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"disable_ap": False, "disable_ble_rpc": False},
+        )
+
+    # Select network
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {CONF_SSID: "MyNetwork"},
+    )
+
+    # Provision - with both disabled, secure device method should not create device
+    with (
+        patch("homeassistant.components.shelly.config_flow.async_provision_wifi"),
+        patch(
+            "homeassistant.components.shelly.config_flow.async_lookup_device_by_name",
+            return_value=("1.1.1.1", 80),
+        ),
+        patch(
+            "homeassistant.components.shelly.config_flow.get_info",
+            return_value=MOCK_DEVICE_INFO,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_PASSWORD: "my_password"},
+        )
+        await hass.async_block_till_done()
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    # Verify entry created (secure device call is skipped when both disabled)
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+
+
+async def test_bluetooth_provision_secure_device_only_ap_disabled(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_setup: AsyncMock,
+) -> None:
+    """Test provisioning with only AP disable enabled."""
+    inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        data=BLE_DISCOVERY_INFO,
+        context={"source": config_entries.SOURCE_BLUETOOTH},
+    )
+
+    # Confirm with only AP disable
+    with patch(
+        "homeassistant.components.shelly.config_flow.async_scan_wifi_networks",
+        return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}],
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"disable_ap": True, "disable_ble_rpc": False},
+        )
+
+    # Select network
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {CONF_SSID: "MyNetwork"},
+    )
+
+    # Provision and verify only AP disabled
+    mock_device = AsyncMock()
+    mock_device.initialize = AsyncMock()
+    mock_device.wifi_setconfig = AsyncMock(return_value={})
+    mock_device.shutdown = AsyncMock()
+
+    with (
+        patch("homeassistant.components.shelly.config_flow.async_provision_wifi"),
+        patch(
+            "homeassistant.components.shelly.config_flow.async_lookup_device_by_name",
+            return_value=("1.1.1.1", 80),
+        ),
+        patch(
+            "homeassistant.components.shelly.config_flow.get_info",
+            return_value=MOCK_DEVICE_INFO,
+        ),
+        patch(
+            "homeassistant.components.shelly.config_flow.RpcDevice.create",
+            return_value=mock_device,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_PASSWORD: "my_password"},
+        )
+        await hass.async_block_till_done()
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    # Verify entry created
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+
+    # Verify only wifi_setconfig was called
+    mock_device.wifi_setconfig.assert_called_once_with(ap_enable=False)
+    assert mock_device.shutdown.called
+
+
+async def test_bluetooth_provision_secure_device_only_ble_disabled(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_setup: AsyncMock,
+) -> None:
+    """Test provisioning with only BLE disable enabled."""
+    inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        data=BLE_DISCOVERY_INFO,
+        context={"source": config_entries.SOURCE_BLUETOOTH},
+    )
+
+    # Confirm with only BLE disable
+    with patch(
+        "homeassistant.components.shelly.config_flow.async_scan_wifi_networks",
+        return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}],
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"disable_ap": False, "disable_ble_rpc": True},
+        )
+
+    # Select network
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {CONF_SSID: "MyNetwork"},
+    )
+
+    # Provision and verify only BLE disabled
+    mock_device = AsyncMock()
+    mock_device.initialize = AsyncMock()
+    mock_device.ble_setconfig = AsyncMock(return_value={"restart_required": False})
+    mock_device.shutdown = AsyncMock()
+
+    with (
+        patch("homeassistant.components.shelly.config_flow.async_provision_wifi"),
+        patch(
+            "homeassistant.components.shelly.config_flow.async_lookup_device_by_name",
+            return_value=("1.1.1.1", 80),
+        ),
+        patch(
+            "homeassistant.components.shelly.config_flow.get_info",
+            return_value=MOCK_DEVICE_INFO,
+        ),
+        patch(
+            "homeassistant.components.shelly.config_flow.RpcDevice.create",
+            return_value=mock_device,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_PASSWORD: "my_password"},
+        )
+        await hass.async_block_till_done()
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    # Verify entry created
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+
+    # Verify only ble_setconfig was called
+    mock_device.ble_setconfig.assert_called_once_with(enable=True, enable_rpc=False)
+    assert mock_device.shutdown.called
+
+
+async def test_bluetooth_provision_secure_device_with_restart_required(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_setup: AsyncMock,
+) -> None:
+    """Test provisioning when BLE disable requires restart."""
+    inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        data=BLE_DISCOVERY_INFO,
+        context={"source": config_entries.SOURCE_BLUETOOTH},
+    )
+
+    # Confirm with both enabled
+    with patch(
+        "homeassistant.components.shelly.config_flow.async_scan_wifi_networks",
+        return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}],
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"disable_ap": True, "disable_ble_rpc": True},
+        )
+
+    # Select network
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {CONF_SSID: "MyNetwork"},
+    )
+
+    # Provision and verify restart is triggered
+    mock_device = AsyncMock()
+    mock_device.initialize = AsyncMock()
+    mock_device.wifi_setconfig = AsyncMock(return_value={})
+    mock_device.ble_setconfig = AsyncMock(return_value={"restart_required": True})
+    mock_device.trigger_reboot = AsyncMock()
+    mock_device.shutdown = AsyncMock()
+
+    with (
+        patch("homeassistant.components.shelly.config_flow.async_provision_wifi"),
+        patch(
+            "homeassistant.components.shelly.config_flow.async_lookup_device_by_name",
+            return_value=("1.1.1.1", 80),
+        ),
+        patch(
+            "homeassistant.components.shelly.config_flow.get_info",
+            return_value=MOCK_DEVICE_INFO,
+        ),
+        patch(
+            "homeassistant.components.shelly.config_flow.RpcDevice.create",
+            return_value=mock_device,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_PASSWORD: "my_password"},
+        )
+        await hass.async_block_till_done()
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    # Verify entry created
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+
+    # Verify restart was triggered and shutdown called
+    mock_device.trigger_reboot.assert_called_once_with(delay_ms=1000)
+    assert mock_device.shutdown.called
+
+
+async def test_bluetooth_provision_secure_device_fails_gracefully(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_setup: AsyncMock,
+) -> None:
+    """Test provisioning succeeds even when secure device calls fail."""
+    inject_bluetooth_service_info_bleak(hass, BLE_DISCOVERY_INFO)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN,
+        data=BLE_DISCOVERY_INFO,
+        context={"source": config_entries.SOURCE_BLUETOOTH},
+    )
+
+    # Confirm with both enabled
+    with patch(
+        "homeassistant.components.shelly.config_flow.async_scan_wifi_networks",
+        return_value=[{"ssid": "MyNetwork", "rssi": -50, "auth": 2}],
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {"disable_ap": True, "disable_ble_rpc": True},
+        )
+
+    # Select network
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {CONF_SSID: "MyNetwork"},
+    )
+
+    # Provision with security calls failing - wifi_setconfig will fail
+    mock_device = AsyncMock()
+    mock_device.initialize = AsyncMock()
+    mock_device.wifi_setconfig = AsyncMock(side_effect=RpcCallError("RPC call failed"))
+    mock_device.shutdown = AsyncMock()
+
+    with (
+        patch("homeassistant.components.shelly.config_flow.async_provision_wifi"),
+        patch(
+            "homeassistant.components.shelly.config_flow.async_lookup_device_by_name",
+            return_value=("1.1.1.1", 80),
+        ),
+        patch(
+            "homeassistant.components.shelly.config_flow.get_info",
+            return_value=MOCK_DEVICE_INFO,
+        ),
+        patch(
+            "homeassistant.components.shelly.config_flow.RpcDevice.create",
+            return_value=mock_device,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {CONF_PASSWORD: "my_password"},
+        )
+        await hass.async_block_till_done()
+        result = await hass.config_entries.flow.async_configure(result["flow_id"])
+
+    # Verify entry still created despite secure device failure
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    assert result["result"].unique_id == "C049EF8873E8"
+
+
 async def test_zeroconf_aborts_idle_ble_flow(
     hass: HomeAssistant,
     mock_setup_entry: AsyncMock,