From ed8a372f4eeecc80aaa1ec2a3265aa4672f28c6d Mon Sep 17 00:00:00 2001
From: TJ Horner <me@tjhorner.com>
Date: Fri, 6 Oct 2023 11:00:04 -0700
Subject: [PATCH] Auto-fix common key entry issues during WeatherKit config
 flow (#101504)

---
 .../components/weatherkit/config_flow.py      | 20 ++++++++
 .../components/weatherkit/test_config_flow.py | 51 +++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/homeassistant/components/weatherkit/config_flow.py b/homeassistant/components/weatherkit/config_flow.py
index 5762c4ae9b2..657a80547ab 100644
--- a/homeassistant/components/weatherkit/config_flow.py
+++ b/homeassistant/components/weatherkit/config_flow.py
@@ -66,6 +66,7 @@ class WeatherKitFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
         errors = {}
         if user_input is not None:
             try:
+                user_input[CONF_KEY_PEM] = self._fix_key_input(user_input[CONF_KEY_PEM])
                 await self._test_config(user_input)
             except WeatherKitUnsupportedLocationError as exception:
                 LOGGER.error(exception)
@@ -104,6 +105,25 @@ class WeatherKitFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
             errors=errors,
         )
 
+    def _fix_key_input(self, key_input: str) -> str:
+        """Fix common user errors with the key input."""
+        # OSes may sometimes turn two hyphens (--) into an em dash (—)
+        key_input = key_input.replace("—", "--")
+
+        # Trim whitespace and line breaks
+        key_input = key_input.strip()
+
+        # Make sure header and footer are present
+        header = "-----BEGIN PRIVATE KEY-----"
+        if not key_input.startswith(header):
+            key_input = f"{header}\n{key_input}"
+
+        footer = "-----END PRIVATE KEY-----"
+        if not key_input.endswith(footer):
+            key_input += f"\n{footer}"
+
+        return key_input
+
     async def _test_config(self, user_input: dict[str, Any]) -> None:
         """Validate credentials."""
         client = WeatherKitApiClient(
diff --git a/tests/components/weatherkit/test_config_flow.py b/tests/components/weatherkit/test_config_flow.py
index 3b6cf76a3d5..9e4d03cbad4 100644
--- a/tests/components/weatherkit/test_config_flow.py
+++ b/tests/components/weatherkit/test_config_flow.py
@@ -126,3 +126,54 @@ async def test_form_unsupported_location(hass: HomeAssistant) -> None:
         )
 
     assert result["type"] == FlowResultType.CREATE_ENTRY
+
+
+@pytest.mark.parametrize(
+    ("input_header"),
+    [
+        "-----BEGIN PRIVATE KEY-----\n",
+        "",
+        "  \n\n-----BEGIN PRIVATE KEY-----\n",
+        "—---BEGIN PRIVATE KEY-----\n",
+    ],
+    ids=["Correct header", "No header", "Leading characters", "Em dash in header"],
+)
+@pytest.mark.parametrize(
+    ("input_footer"),
+    [
+        "\n-----END PRIVATE KEY-----",
+        "",
+        "\n-----END PRIVATE KEY-----\n\n  ",
+        "\n—---END PRIVATE KEY-----",
+    ],
+    ids=["Correct footer", "No footer", "Trailing characters", "Em dash in footer"],
+)
+async def test_auto_fix_key_input(
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    input_header: str,
+    input_footer: str,
+) -> None:
+    """Test that we fix common user errors in key input."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == FlowResultType.FORM
+    assert result["errors"] == {}
+
+    with patch(
+        "homeassistant.components.weatherkit.WeatherKitApiClient.get_availability",
+        return_value=[DataSetType.CURRENT_WEATHER],
+    ):
+        user_input = EXAMPLE_USER_INPUT.copy()
+        user_input[CONF_KEY_PEM] = f"{input_header}whateverkey{input_footer}"
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            user_input,
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == FlowResultType.CREATE_ENTRY
+
+    assert result["data"][CONF_KEY_PEM] == EXAMPLE_CONFIG_DATA[CONF_KEY_PEM]
+    assert len(mock_setup_entry.mock_calls) == 1