From a0ff08436ed6005d7b55fff4a00867690fe31522 Mon Sep 17 00:00:00 2001 From: Cherry Date: Thu, 25 Jun 2026 10:14:01 +1200 Subject: [PATCH] Fix China mainland auth endpoints and 2FA fallback --- pyicloud/base.py | 9 ++++++++- pyicloud/cli/context.py | 8 ++++++-- tests/test_base.py | 27 +++++++++++++++++++++++++++ tests/test_cmdline.py | 20 ++++++++++++++------ 4 files changed, 55 insertions(+), 9 deletions(-) diff --git a/pyicloud/base.py b/pyicloud/base.py index c63b10b3..cb8ed1a7 100644 --- a/pyicloud/base.py +++ b/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.""" diff --git a/pyicloud/cli/context.py b/pyicloud/cli/context.py index 84d29791..0add554b 100644 --- a/pyicloud/cli/context.py +++ b/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") diff --git a/tests/test_base.py b/tests/test_base.py index c4aadcc0..c1b2929a 100644 --- a/tests/test_base.py +++ b/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 ( diff --git a/tests/test_cmdline.py b/tests/test_cmdline.py index 8984e3dd..e8f40644 100644 --- a/tests/test_cmdline.py +++ b/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,8 +2657,8 @@ 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 @@ -2661,11 +2666,14 @@ def test_sms_2fa_request_failure_aborts() -> None: "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: