Allow Apple TV pairing PIN to accept leading zeroes (#179258)

Co-authored-by: Josef Zweck <josef@zweck.dev>
Signed-off-by: Sankalp Thakur <sankalphimself@gmail.com>

homeassistant/components/apple_tv/config_flow.py

@@ -28,6 +28,7 @@ from homeassistant.const import CONF_ADDRESS, CONF_NAME, CONF_PIN
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.data_entry_flow import AbortFlow
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.schema_config_entry_flow import (
     SchemaFlowFormStep,
@@ -41,7 +42,7 @@ _LOGGER = logging.getLogger(__name__)
 
 DEVICE_INPUT = "device_input"
 
-INPUT_PIN_SCHEMA = vol.Schema({vol.Required(CONF_PIN, default=None): int})
+INPUT_PIN_SCHEMA = vol.Schema({vol.Required(CONF_PIN, default=""): cv.string})
 
 DEFAULT_START_OFF = False
 
@@ -513,17 +514,23 @@ class AppleTVConfigFlow(ConfigFlow, domain=DOMAIN):
         assert self.pairing
         assert self.protocol
         if user_input is not None:
-            try:
-                self.pairing.pin(user_input[CONF_PIN])
-                await self.pairing.finish()
-                self.credentials[self.protocol.value] = self.pairing.service.credentials
-                return await self.async_pair_next_protocol()
-            except exceptions.PairingError:
-                _LOGGER.exception("Authentication problem")
-                errors["base"] = "invalid_auth"
-            except Exception:
-                _LOGGER.exception("Unexpected exception")
-                errors["base"] = "unknown"
+            pin = user_input[CONF_PIN]
+            if not pin.isascii() or not pin.isdigit():
+                errors["pin"] = "invalid_pin"
+            else:
+                try:
+                    self.pairing.pin(pin)
+                    await self.pairing.finish()
+                    self.credentials[self.protocol.value] = (
+                        self.pairing.service.credentials
+                    )
+                    return await self.async_pair_next_protocol()
+                except exceptions.PairingError:
+                    _LOGGER.exception("Authentication problem")
+                    errors["base"] = "invalid_auth"
+                except Exception:
+                    _LOGGER.exception("Unexpected exception")
+                    errors["base"] = "unknown"
 
         return self.async_show_form(
             step_id="pair_with_pin",

homeassistant/components/apple_tv/strings.json

@@ -17,6 +17,7 @@
     "error": {
       "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
       "invalid_auth": "[%key:common::config_flow::error::invalid_auth%]",
+      "invalid_pin": "Invalid PIN",
       "no_devices_found": "[%key:common::config_flow::abort::no_devices_found%]",
       "unknown": "[%key:common::config_flow::error::unknown%]"
     },
@@ -34,7 +35,7 @@
         "data": {
           "pin": "[%key:common::config_flow::data::pin%]"
         },
-        "description": "Pairing is required for the `{protocol}` protocol. Please enter the PIN code displayed on screen. Leading zeros shall be omitted, i.e. enter 123 if the displayed code is 0123.",
+        "description": "Pairing is required for the `{protocol}` protocol. Please enter the PIN code displayed on screen, including any leading zeros.",
         "title": "Pairing"
       },
       "password": {

tests/components/apple_tv/conftest.py

@@ -61,9 +61,11 @@ def pairing() -> Generator[AsyncMock]:
                 await http.create_session(session), config.get_service(protocol)
             )
             handler.always_fail = mock_pair.always_fail
+            mock_pair.handler = handler
             return handler
 
         mock_pair.always_fail = False
+        mock_pair.handler = None
         mock_pair.side_effect = _pair
         yield mock_pair
 

tests/components/apple_tv/test_config_flow.py

@@ -155,6 +155,59 @@ async def test_user_adds_full_device(hass: HomeAssistant) -> None:
     }
 
 
+@pytest.mark.usefixtures("mrp_device")
+async def test_user_pair_leading_zero_pin(
+    hass: HomeAssistant, pairing: AsyncMock
+) -> None:
+    """Test that a pairing PIN with a leading zero is passed through as a string."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {"device_input": "MRP Device"},
+    )
+    result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "pair_with_pin"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"pin": "0123"}
+    )
+    assert pairing.handler.pin_code == "0123"
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+
+
+@pytest.mark.usefixtures("mrp_device")
+@pytest.mark.parametrize("invalid_pin", ["abcd", "12ab", "١٢٣٤", "123\n"])
+async def test_user_pair_non_numeric_pin(
+    hass: HomeAssistant, pairing: AsyncMock, invalid_pin: str
+) -> None:
+    """Test that a non-numeric PIN is rejected at the form."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {"device_input": "MRP Device"},
+    )
+    result = await hass.config_entries.flow.async_configure(result["flow_id"], {})
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "pair_with_pin"
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"pin": invalid_pin}
+    )
+    assert result["type"] is FlowResultType.FORM
+    assert result["errors"] == {"pin": "invalid_pin"}
+
+    result = await hass.config_entries.flow.async_configure(
+        result["flow_id"], {"pin": "0123"}
+    )
+    assert pairing.handler.pin_code == "0123"
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+
+
 @pytest.mark.usefixtures("dmap_device", "dmap_pin", "pairing")
 async def test_user_adds_dmap_device(hass: HomeAssistant) -> None:
     """Test adding device with only DMAP service."""