fix(security): apply fixes for URL injection and redirect URI validation (#49)

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: deviantintegral <255023+deviantintegral@users.noreply.github.com>

Author: Copilot

src/flameconnect/b2c_login.py

@@ -13,10 +13,13 @@
 import aiohttp
 import yarl
 
+from flameconnect.const import CLIENT_ID
 from flameconnect.exceptions import AuthenticationError
 
 _LOGGER = logging.getLogger(__name__)
 
+_REDIRECT_URI_PREFIX = f"msal{CLIENT_ID}://auth"
+
 
 _B2C_POLICY = "B2C_1A_FirePhoneSignUpOrSignInWithPhoneOrEmail"
 
@@ -287,8 +290,11 @@ async def b2c_login_with_credentials(auth_uri: str, email: str, password: str) -
                                 raise AuthenticationError(
                                     "Redirect without Location header"
                                 )
-                            # Custom-scheme redirect (msal{id}://auth)
-                            if location.startswith("msal") and "://auth" in location:
+                            # Custom-scheme redirect (msal{CLIENT_ID}://auth)
+                            if location.startswith(_REDIRECT_URI_PREFIX) and (
+                                len(location) == len(_REDIRECT_URI_PREFIX)
+                                or location[len(_REDIRECT_URI_PREFIX)] in ("?", "/")
+                            ):
                                 _LOGGER.debug(
                                     "Captured custom-scheme redirect: %s",
                                     location[:120] + "...",

src/flameconnect/client.py

@@ -6,6 +6,7 @@
 import logging
 from dataclasses import replace
 from typing import TYPE_CHECKING, Any
+from urllib.parse import quote
 
 import aiohttp
 
@@ -210,7 +211,7 @@ async def get_fire_overview(self, fire_id: str) -> FireOverview:
         Returns:
             A FireOverview containing the fire identity and decoded parameters.
         """
-        url = f"{API_BASE}/api/Fires/GetFireOverview?FireId={fire_id}"
+        url = f"{API_BASE}/api/Fires/GetFireOverview?FireId={quote(fire_id, safe='')}"
         data: dict[str, Any] = await self._request("GET", url)
 
         wifi: dict[str, Any] = data["WifiFireOverview"]

tests/test_b2c_login.py

@@ -2341,3 +2341,145 @@ async def test_no_none_in_log_method_or_url(self, caplog):
                 assert " None" not in line.split(">>>")[1], (
                     f"Found 'None' in log line: {line}"
                 )
+
+
+# -------------------------------------------------------------------
+# Security regression tests
+# -------------------------------------------------------------------
+
+
+class TestRedirectUriValidation:
+    """Security tests: redirect URI must match the exact expected prefix.
+
+    Before the fix, the check was:
+        location.startswith("msal") and "://auth" in location
+
+    This was too broad. A compromised server could send:
+      - msal{CLIENT_ID}://auth.evil.com?code=...  (domain mismatch)
+      - msalEVIL://auth?code=...                  (wrong client ID)
+
+    Both would match the old check but must NOT match the new one,
+    which requires startswith(f"msal{CLIENT_ID}://auth") followed by
+    '?' or '/'.
+    """
+
+    async def test_correct_redirect_uri_accepted(self):
+        """The legitimate redirect URI is still accepted."""
+        login_resp = _make_mock_response(
+            status=200,
+            text=SAMPLE_B2C_HTML,
+            url=SAMPLE_PAGE_URL,
+        )
+        post_resp = _make_mock_response(status=200, text='{"status":"200"}')
+        confirmed_resp = _make_mock_response(
+            status=302,
+            headers={"Location": REDIRECT_URL},
+        )
+
+        session = _make_mock_session(
+            get=MagicMock(side_effect=[login_resp, confirmed_resp]),
+            post=MagicMock(return_value=post_resp),
+        )
+
+        with _patch_session(session):
+            result = await b2c_login_with_credentials(AUTH_URI, "user@test.com", "pass")
+
+        assert result == REDIRECT_URL
+
+    async def test_correct_redirect_uri_with_trailing_slash_accepted(self):
+        """msal{CLIENT_ID}://auth/?... (with trailing slash) is also valid."""
+        redirect_with_slash = (
+            f"msal{_CLIENT_ID}://auth/?code=test-auth-code-123&state=test-state"
+        )
+        login_resp = _make_mock_response(
+            status=200,
+            text=SAMPLE_B2C_HTML,
+            url=SAMPLE_PAGE_URL,
+        )
+        post_resp = _make_mock_response(status=200, text='{"status":"200"}')
+        confirmed_resp = _make_mock_response(
+            status=302,
+            headers={"Location": redirect_with_slash},
+        )
+
+        session = _make_mock_session(
+            get=MagicMock(side_effect=[login_resp, confirmed_resp]),
+            post=MagicMock(return_value=post_resp),
+        )
+
+        with _patch_session(session):
+            result = await b2c_login_with_credentials(AUTH_URI, "user@test.com", "pass")
+
+        assert result == redirect_with_slash
+
+    async def test_attacker_subdomain_redirect_not_accepted(self):
+        """msal{CLIENT_ID}://auth.evil.com?... must NOT be accepted.
+
+        The old check '://auth' in location would accept this because
+        the substring '://auth' appears in '://auth.evil.com'.
+        The new check requires startswith(f'msal{CLIENT_ID}://auth')
+        followed by '?' or '/', which the malicious URL fails because
+        '.' is neither.
+        """
+        malicious_redirect = f"msal{_CLIENT_ID}://auth.evil.com?code=stolen&state=xyz"
+        login_resp = _make_mock_response(
+            status=200,
+            text=SAMPLE_B2C_HTML,
+            url=SAMPLE_PAGE_URL,
+        )
+        post_resp = _make_mock_response(status=200, text='{"status":"200"}')
+        # The malicious redirect followed by a 200 with no real redirect
+        malicious_resp = _make_mock_response(
+            status=302,
+            headers={"Location": malicious_redirect},
+        )
+        final_200 = _make_mock_response(
+            status=200,
+            text="<html>no redirect here</html>",
+        )
+
+        session = _make_mock_session(
+            get=MagicMock(side_effect=[login_resp, malicious_resp, final_200]),
+            post=MagicMock(return_value=post_resp),
+        )
+
+        with (
+            _patch_session(session),
+            pytest.raises(
+                AuthenticationError,
+                match="Reached 200 response without finding redirect URL",
+            ),
+        ):
+            await b2c_login_with_credentials(AUTH_URI, "user@test.com", "pass")
+
+    async def test_wrong_client_id_redirect_not_accepted(self):
+        """msalWRONG_CLIENT_ID://auth?... must NOT be accepted."""
+        malicious_redirect = "msalDEADBEEF-0000-0000-0000-000000000000://auth?code=x"
+        login_resp = _make_mock_response(
+            status=200,
+            text=SAMPLE_B2C_HTML,
+            url=SAMPLE_PAGE_URL,
+        )
+        post_resp = _make_mock_response(status=200, text='{"status":"200"}')
+        malicious_resp = _make_mock_response(
+            status=302,
+            headers={"Location": malicious_redirect},
+        )
+        final_200 = _make_mock_response(
+            status=200,
+            text="<html>no redirect here</html>",
+        )
+
+        session = _make_mock_session(
+            get=MagicMock(side_effect=[login_resp, malicious_resp, final_200]),
+            post=MagicMock(return_value=post_resp),
+        )
+
+        with (
+            _patch_session(session),
+            pytest.raises(
+                AuthenticationError,
+                match="Reached 200 response without finding redirect URL",
+            ),
+        ):
+            await b2c_login_with_credentials(AUTH_URI, "user@test.com", "pass")

tests/test_client.py

@@ -1419,3 +1419,56 @@ async def test_turn_off_no_mode_uses_default_temp(self, mock_api, token_auth):
         assert raw[3] == FireMode.STANDBY
         temp = float(raw[4]) + float(raw[5]) / 10.0
         assert temp == pytest.approx(22.0)
+
+
+# -------------------------------------------------------------------
+# Security regression tests
+# -------------------------------------------------------------------
+
+
+class TestGetFireOverviewUrlEncoding:
+    """Security tests: fire_id must be URL-encoded in get_fire_overview.
+
+    A compromised API server could return a FireId containing URL
+    special characters (e.g. '&evil=true') to inject extra query
+    parameters into subsequent requests. The fix URL-encodes fire_id
+    via urllib.parse.quote() so injected characters are escaped.
+    """
+
+    async def test_ampersand_in_fire_id_is_encoded(self, mock_api, token_auth):
+        """'&' in fire_id must be percent-encoded, not passed raw."""
+        fire_id = "abc&evil=true"
+        # The URL must encode '&' as '%26' and '=' as '%3D'
+        encoded_id = "abc%26evil%3Dtrue"
+        url = f"{API_BASE}/api/Fires/GetFireOverview?FireId={encoded_id}"
+        mock_api.get(url, payload={"WifiFireOverview": {"FireId": fire_id}})
+
+        async with FlameConnectClient(token_auth) as client:
+            overview = await client.get_fire_overview(fire_id)
+
+        assert overview.fire.fire_id == fire_id
+
+    async def test_hash_in_fire_id_is_encoded(self, mock_api, token_auth):
+        """'#' in fire_id must be percent-encoded."""
+        fire_id = "abc#fragment"
+        encoded_id = "abc%23fragment"
+        url = f"{API_BASE}/api/Fires/GetFireOverview?FireId={encoded_id}"
+        mock_api.get(url, payload={"WifiFireOverview": {"FireId": fire_id}})
+
+        async with FlameConnectClient(token_auth) as client:
+            overview = await client.get_fire_overview(fire_id)
+
+        assert overview.fire.fire_id == fire_id
+
+    async def test_normal_fire_id_unchanged(
+        self, mock_api, token_auth, get_fire_overview_payload
+    ):
+        """A plain alphanumeric fire_id is unaffected by encoding."""
+        fire_id = "test-fire-001"
+        url = f"{API_BASE}/api/Fires/GetFireOverview?FireId={fire_id}"
+        mock_api.get(url, payload=get_fire_overview_payload)
+
+        async with FlameConnectClient(token_auth) as client:
+            overview = await client.get_fire_overview(fire_id)
+
+        assert overview.fire.fire_id == fire_id