From ecc0917e8f40bcf5a8f836d1d422975bf3f70cae Mon Sep 17 00:00:00 2001
From: puddly <32534428+puddly@users.noreply.github.com>
Date: Thu, 6 Jul 2023 11:47:51 -0400
Subject: [PATCH] Migrate bracketed IP addresses in ZHA config entry (#95917)

* Automatically correct IP addresses surrounded by brackets

* Simplify regex

* Move pattern inline

* Maintain old behavior of stripping whitespace
---
 homeassistant/components/zha/__init__.py | 24 ++++++++++++++++++++----
 tests/components/zha/test_init.py        | 12 ++++++++++--
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/zha/__init__.py b/homeassistant/components/zha/__init__.py
index 5607cabffea..8a81648b580 100644
--- a/homeassistant/components/zha/__init__.py
+++ b/homeassistant/components/zha/__init__.py
@@ -3,6 +3,7 @@ import asyncio
 import copy
 import logging
 import os
+import re
 
 import voluptuous as vol
 from zhaquirks import setup as setup_quirks
@@ -85,19 +86,34 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
+def _clean_serial_port_path(path: str) -> str:
+    """Clean the serial port path, applying corrections where necessary."""
+
+    if path.startswith("socket://"):
+        path = path.strip()
+
+    # Removes extraneous brackets from IP addresses (they don't parse in CPython 3.11.4)
+    if re.match(r"^socket://\[\d+\.\d+\.\d+\.\d+\]:\d+$", path):
+        path = path.replace("[", "").replace("]", "")
+
+    return path
+
+
 async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Set up ZHA.
 
     Will automatically load components to support devices found on the network.
     """
 
-    # Strip whitespace around `socket://` URIs, this is no longer accepted by zigpy
-    # This will be removed in 2023.7.0
+    # Remove brackets around IP addresses, this no longer works in CPython 3.11.4
+    # This will be removed in 2023.11.0
     path = config_entry.data[CONF_DEVICE][CONF_DEVICE_PATH]
+    cleaned_path = _clean_serial_port_path(path)
     data = copy.deepcopy(dict(config_entry.data))
 
-    if path.startswith("socket://") and path != path.strip():
-        data[CONF_DEVICE][CONF_DEVICE_PATH] = path.strip()
+    if path != cleaned_path:
+        _LOGGER.debug("Cleaned serial port path %r -> %r", path, cleaned_path)
+        data[CONF_DEVICE][CONF_DEVICE_PATH] = cleaned_path
         hass.config_entries.async_update_entry(config_entry, data=data)
 
     zha_data = hass.data.setdefault(DATA_ZHA, {})
diff --git a/tests/components/zha/test_init.py b/tests/components/zha/test_init.py
index 23a76de4c25..24ee63fb3d5 100644
--- a/tests/components/zha/test_init.py
+++ b/tests/components/zha/test_init.py
@@ -114,19 +114,27 @@ async def test_config_depreciation(hass: HomeAssistant, zha_config) -> None:
 @pytest.mark.parametrize(
     ("path", "cleaned_path"),
     [
+        # No corrections
         ("/dev/path1", "/dev/path1"),
+        ("/dev/path1[asd]", "/dev/path1[asd]"),
         ("/dev/path1 ", "/dev/path1 "),
+        ("socket://1.2.3.4:5678", "socket://1.2.3.4:5678"),
+        # Brackets around URI
+        ("socket://[1.2.3.4]:5678", "socket://1.2.3.4:5678"),
+        # Spaces
         ("socket://dev/path1 ", "socket://dev/path1"),
+        # Both
+        ("socket://[1.2.3.4]:5678 ", "socket://1.2.3.4:5678"),
     ],
 )
 @patch("homeassistant.components.zha.setup_quirks", Mock(return_value=True))
 @patch(
     "homeassistant.components.zha.websocket_api.async_load_api", Mock(return_value=True)
 )
-async def test_setup_with_v3_spaces_in_uri(
+async def test_setup_with_v3_cleaning_uri(
     hass: HomeAssistant, path: str, cleaned_path: str
 ) -> None:
-    """Test migration of config entry from v3 with spaces after `socket://` URI."""
+    """Test migration of config entry from v3, applying corrections to the port path."""
     config_entry_v3 = MockConfigEntry(
         domain=DOMAIN,
         data={