Fix China mainland auth endpoints and 2FA fallback

pyicloud/base.py

@@ -218,7 +218,7 @@ def _setup_endpoints(self) -> None:
         # If the country or region setting of your Apple ID is China mainland.
         # See https://support.apple.com/en-us/HT208351
         icloud_china: str = ".cn" if self._is_china_mainland else ""
-        self._idmsa_endpoint: str = f"https://idmsa.apple.com{icloud_china}"
+        self._idmsa_endpoint: str = "https://idmsa.apple.com"
         self._auth_endpoint: str = f"{self._idmsa_endpoint}/appleauth/auth"
         self._home_endpoint: str = f"https://www.icloud.com{icloud_china}"
         self._setup_endpoint: str = f"https://setup.icloud.com{icloud_china}/setup/ws/1"
@@ -711,6 +711,7 @@ def _get_auth_headers(
         headers.update(
             {
                 "Referer": self._idmsa_endpoint,
+                "X-Apple-OAuth-Redirect-URI": self._home_endpoint,
                 "X-Apple-OAuth-State": self._client_id,
                 "X-Apple-Frame-Id": self._client_id,
             }
@@ -871,6 +872,12 @@ def _set_two_factor_delivery_state(
         self._two_factor_delivery_method = method
         self._two_factor_delivery_notice = notice
 
+    def use_existing_trusted_device_code(self) -> None:
+        """Validate the next 2FA code as one already shown on a trusted device."""
+
+        self._clear_trusted_device_bridge_state()
+        self._set_two_factor_delivery_state("trusted_device")
+
     def _current_hsa2_boot_context(self) -> Hsa2BootContext:
         """Return the best available HSA2 boot context for the active challenge."""
 

pyicloud/cli/context.py

@@ -368,8 +368,12 @@ def _handle_2fa(self, api: PyiCloudService) -> None:
                 raise CLIAbort(
                     "Failed to request the 2FA trusted-device prompt."
                 ) from exc
-            except PyiCloudAPIResponseException as exc:
-                raise CLIAbort("Failed to request the 2FA SMS code.") from exc
+            except PyiCloudAPIResponseException:
+                self.console.print(
+                    "Failed to request the 2FA SMS code. "
+                    "If your trusted Apple device already shows a code, enter it below."
+                )
+                api.use_existing_trusted_device_code()
             max_attempts = 3
             for attempt in range(max_attempts):
                 code = typer.prompt("Enter 2FA code")

tests/test_base.py

@@ -111,6 +111,33 @@ def test_constructor_skips_authentication_when_requested() -> None:
         get_from_keyring.assert_not_called()
 
 
+def test_china_mainland_uses_global_idmsa_and_cn_icloud_endpoints() -> None:
+    """China mainland accounts use global IDMS auth and China iCloud services."""
+    with (
+        patch("pyicloud.PyiCloudService.authenticate") as mock_authenticate,
+        patch("pyicloud.PyiCloudService._setup_cookie_directory") as mock_setup_dir,
+        patch("builtins.open", new_callable=mock_open),
+    ):
+        mock_setup_dir.return_value = "/tmp/pyicloud/cookies"
+
+        service = PyiCloudService(
+            "test@example.com",
+            secrets.token_hex(32),
+            china_mainland=True,
+            authenticate=False,
+        )
+
+    assert service._idmsa_endpoint == "https://idmsa.apple.com"
+    assert service._auth_endpoint == "https://idmsa.apple.com/appleauth/auth"
+    assert service._home_endpoint == "https://www.icloud.com.cn"
+    assert service._setup_endpoint == "https://setup.icloud.com.cn/setup/ws/1"
+    assert (
+        service._get_auth_headers()["X-Apple-OAuth-Redirect-URI"]
+        == "https://www.icloud.com.cn"
+    )
+    mock_authenticate.assert_not_called()
+
+
 def test_constructor_accepts_keyword_only_cloudkit_validation_extra() -> None:
     """cloudkit_validation_extra remains a keyword-only escape hatch."""
     with (

tests/test_cmdline.py

@@ -1011,6 +1011,11 @@ def __init__(
         self.two_factor_delivery_method = "unknown"
         self.two_factor_delivery_notice = None
         self.request_2fa_code = MagicMock(return_value=False)
+        self.use_existing_trusted_device_code = MagicMock(
+            side_effect=lambda: setattr(
+                self, "two_factor_delivery_method", "trusted_device"
+            )
+        )
         self.validate_2fa_code = MagicMock(return_value=True)
         self.confirm_security_key = MagicMock(return_value=True)
         self.send_verification_code = MagicMock(return_value=True)
@@ -2652,20 +2657,23 @@ def request_sms_fallback() -> bool:
     fake_api.validate_2fa_code.assert_called_once_with("123456")
 
 
-def test_sms_2fa_request_failure_aborts() -> None:
-    """Auth login should surface SMS delivery request failures clearly."""
+def test_sms_2fa_request_failure_still_prompts_for_existing_device_code() -> None:
+    """Auth login should accept a code that already appeared on a trusted device."""
 
     fake_api = FakeAPI()
     fake_api.requires_2fa = True
     fake_api.request_2fa_code.side_effect = context_module.PyiCloudAPIResponseException(
         "sms request failed"
     )
 
-    result = _invoke(fake_api, "auth", "login", interactive=True)
+    with patch.object(context_module.typer, "prompt", return_value="123456"):
+        result = _invoke(fake_api, "auth", "login", interactive=True)
 
-    assert result.exit_code != 0
-    assert result.exception.args[0] == "Failed to request the 2FA SMS code."
-    fake_api.validate_2fa_code.assert_not_called()
+    assert result.exit_code == 0
+    assert "Failed to request the 2FA SMS code." in result.stdout
+    fake_api.use_existing_trusted_device_code.assert_called_once_with()
+    assert fake_api.two_factor_delivery_method == "trusted_device"
+    fake_api.validate_2fa_code.assert_called_once_with("123456")
 
 
 def test_trusted_device_2fa_request_failure_aborts() -> None: