From 4adbfb32ec54f8a7838e2c90c690f06f33d71705 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Niccolo=CC=80=20Zapponi?= Date: Thu, 29 Oct 2020 16:51:08 +0000 Subject: [PATCH] Added tests --- pyicloud/base.py | 63 ++++++++++++++++++++----------------------- pyicloud/cmdline.py | 17 +++++++++++- tests/__init__.py | 55 ++++++++++++++++++++++++++++--------- tests/const.py | 9 +++++-- tests/const_login.py | 6 ++++- tests/test_cmdline.py | 8 +++--- 6 files changed, 104 insertions(+), 54 deletions(-) diff --git a/pyicloud/base.py b/pyicloud/base.py index 3f907d9..349cd7e 100644 --- a/pyicloud/base.py +++ b/pyicloud/base.py @@ -101,15 +101,13 @@ class PyiCloudSession(Session): ) # Save session_data to file - with open(self.service._get_sessiondata_path(), "w") as outfile: + with open(self.service.session_path, "w") as outfile: json.dump(self.service.session_data, outfile) LOGGER.debug("Saved session data to file") # Save cookies to file - if not path.exists(self.service._cookie_directory): - mkdir(self.service._cookie_directory) self.cookies.save(ignore_discard=True, ignore_expires=True) - LOGGER.debug("Cookies saved to %s", self.service._get_cookiejar_path()) + LOGGER.debug("Cookies saved to %s", self.service.cookiejar_path) if not response.ok and content_type not in json_mimetypes: if has_retried is None and response.status_code == 450: @@ -214,16 +212,14 @@ class PyiCloudService(object): if session_directory: self._session_directory = session_directory else: - self._session_directory = path.join( - gettempdir(), "pyicloud-session" - ) - LOGGER.debug(f"Using session file {self._get_sessiondata_path()}") + self._session_directory = path.join(gettempdir(), "pyicloud-session") + LOGGER.debug("Using session file %s", self.session_path) try: - with open(self._get_sessiondata_path()) as session_f: + with open(self.session_path) as session_f: self.session_data = json.load(session_f) - except: - LOGGER.warning("Session file does not exist") + except: # pylint: disable=bare-except + LOGGER.info("Session file does not exist") if not path.exists(self._session_directory): mkdir(self._session_directory) @@ -236,6 +232,9 @@ class PyiCloudService(object): else: self._cookie_directory = path.join(gettempdir(), "pyicloud") + if not path.exists(self._cookie_directory): + mkdir(self._cookie_directory) + if self.session_data.get("client_id"): self.client_id = self.session_data.get("client_id") @@ -245,7 +244,7 @@ class PyiCloudService(object): {"Origin": self.HOME_ENDPOINT, "Referer": f"{self.HOME_ENDPOINT}/"} ) - cookiejar_path = self._get_cookiejar_path() + cookiejar_path = self.cookiejar_path self.session.cookies = cookielib.LWPCookieJar(filename=cookiejar_path) if path.exists(cookiejar_path): try: @@ -277,7 +276,7 @@ class PyiCloudService(object): LOGGER.info("Session token is still valid") self.data = req.json() login_successful = True - except: + except PyiCloudAPIResponseException: msg = "Invalid authentication token, will log in from scratch." if not login_successful: @@ -344,15 +343,17 @@ class PyiCloudService(object): raise PyiCloudFailedLoginException(msg, error) self.data = req.json() - - def _get_cookiejar_path(self): + + @property + def cookiejar_path(self): """Get path for cookiejar file.""" return path.join( self._cookie_directory, "".join([c for c in self.user.get("accountName") if match(r"\w", c)]), ) - def _get_sessiondata_path(self): + @property + def session_path(self): """Get path for session data file.""" return path.join( self._session_directory, @@ -362,23 +363,15 @@ class PyiCloudService(object): @property def requires_2sa(self): """Returns True if two-step authentication is required.""" - return ( - self.data["dsInfo"].get("hsaVersion", 0) >= 1 - and ( - self.data.get("hsaChallengeRequired", False) - or not self.is_trusted_session - ) + return self.data.get("dsInfo", {}).get("hsaVersion", 0) >= 1 and ( + self.data.get("hsaChallengeRequired", False) or not self.is_trusted_session ) @property def requires_2fa(self): """Returns True if two-factor authentication is required.""" - return ( - self.data["dsInfo"].get("hsaVersion", 0) == 2 - and ( - self.data.get("hsaChallengeRequired", False) - or not self.is_trusted_session - ) + return self.data["dsInfo"].get("hsaVersion", 0) == 2 and ( + self.data.get("hsaChallengeRequired", False) or not self.is_trusted_session ) @property @@ -451,19 +444,21 @@ class PyiCloudService(object): headers["X-Apple-ID-Session-Id"] = self.session_data.get("session_id") try: - req = self.session.post( + self.session.post( f"{self.AUTH_ENDPOINT}/verify/trusteddevice/securitycode", data=json.dumps(data), headers=headers, ) except PyiCloudAPIResponseException as error: - LOGGER.error("Code verification failed.") - return False + if error.code == -21669: + # Wrong verification code + LOGGER.error("Code verification failed.") + return False + raise LOGGER.debug("Code verification successful.") self.trust_session() - return not self.requires_2sa def trust_session(self): @@ -487,13 +482,13 @@ class PyiCloudService(object): headers["X-Apple-ID-Session-Id"] = self.session_data.get("session_id") try: - req = self.session.get( + self.session.get( f"{self.AUTH_ENDPOINT}/2sv/trust", headers=headers, ) self._authenticate_with_token() return True - except PyiCloudAPIResponseException as error: + except PyiCloudAPIResponseException: LOGGER.error("Session trust failed.") return False diff --git a/pyicloud/cmdline.py b/pyicloud/cmdline.py index a6bbcf7..8137cda 100644 --- a/pyicloud/cmdline.py +++ b/pyicloud/cmdline.py @@ -201,7 +201,22 @@ def main(args=None): ): utils.store_password_in_keyring(username, password) - if api.requires_2sa: + if api.requires_2fa: + # fmt: off + print( + "\nTwo-step authentication required.", + "\nPlease enter validation code" + ) + # fmt: on + + code = input("(string) --> ") + if not api.validate_2fa_code(code): + print("Failed to verify verification code") + sys.exit(1) + + print("") + + elif api.requires_2sa: # fmt: off print( "\nTwo-step authentication required.", diff --git a/tests/__init__.py b/tests/__init__.py index 828a5dd..d36a28d 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -9,13 +9,19 @@ from pyicloud.services.findmyiphone import FindMyiPhoneServiceManager, AppleDevi from .const import ( AUTHENTICATED_USER, - REQUIRES_2SA_USER, + REQUIRES_2FA_USER, + REQUIRES_2FA_TOKEN, + VALID_TOKEN, VALID_USERS, VALID_PASSWORD, + VALID_COOKIE, + VALID_2FA_CODE, + VALID_TOKENS, ) from .const_login import ( + AUTH_OK, LOGIN_WORKING, - LOGIN_2SA, + LOGIN_2FA, TRUSTED_DEVICES, TRUSTED_DEVICE_1, VERIFICATION_CODE_OK, @@ -41,6 +47,7 @@ class ResponseMock(Response): self.result = result self.status_code = status_code self.raw = kwargs.get("raw") + self.headers = kwargs.get("headers", {}) @property def text(self): @@ -52,21 +59,16 @@ class PyiCloudSessionMock(base.PyiCloudSession): def request(self, method, url, **kwargs): params = kwargs.get("params") + headers = kwargs.get("headers") data = json.loads(kwargs.get("data", "{}")) # Login if self.service.SETUP_ENDPOINT in url: - if "login" in url and method == "POST": - if ( - data.get("apple_id") not in VALID_USERS - or data.get("password") != VALID_PASSWORD - ): + if "accountLogin" in url and method == "POST": + if data.get("dsWebAuthToken") not in VALID_TOKENS: self._raise_error(None, "Unknown reason") - if ( - data.get("apple_id") == REQUIRES_2SA_USER - and data.get("password") == VALID_PASSWORD - ): - return ResponseMock(LOGIN_2SA) + if data.get("dsWebAuthToken") == REQUIRES_2FA_TOKEN: + return ResponseMock(LOGIN_2FA) return ResponseMock(LOGIN_WORKING) if "listDevices" in url and method == "GET": @@ -84,6 +86,35 @@ class PyiCloudSessionMock(base.PyiCloudSession): return ResponseMock(VERIFICATION_CODE_OK) self._raise_error(None, "FOUND_CODE") + if "validate" in url and method == "POST": + if headers.get("X-APPLE-WEBAUTH-TOKEN") == VALID_COOKIE: + return ResponseMock(LOGIN_WORKING) + self._raise_error(None, "Session expired") + + if self.service.AUTH_ENDPOINT in url: + if "signin" in url and method == "POST": + if ( + data.get("accountName") not in VALID_USERS + or data.get("password") != VALID_PASSWORD + ): + self._raise_error(None, "Unknown reason") + if data.get("accountName") == REQUIRES_2FA_USER: + self.service.session_data["session_token"] = REQUIRES_2FA_TOKEN + return ResponseMock(AUTH_OK) + + self.service.session_data["session_token"] = VALID_TOKEN + return ResponseMock(AUTH_OK) + + if "securitycode" in url and method == "POST": + if data.get("securityCode", {}).get("code") != VALID_2FA_CODE: + self._raise_error(None, "Incorrect code") + + self.service.session_data["session_token"] = VALID_TOKEN + return ResponseMock("", status_code=204) + + if "trust" in url and method == "GET": + return ResponseMock("", status_code=204) + # Account if "device/getDevices" in url and method == "GET": return ResponseMock(ACCOUNT_DEVICES_WORKING) diff --git a/tests/const.py b/tests/const.py index 3b69842..1dbd27c 100644 --- a/tests/const.py +++ b/tests/const.py @@ -4,8 +4,13 @@ from .const_account_family import PRIMARY_EMAIL, APPLE_ID_EMAIL, ICLOUD_ID_EMAIL # Base AUTHENTICATED_USER = PRIMARY_EMAIL -REQUIRES_2SA_USER = "requires_2sa_user" -VALID_USERS = [AUTHENTICATED_USER, REQUIRES_2SA_USER, APPLE_ID_EMAIL, ICLOUD_ID_EMAIL] +REQUIRES_2FA_TOKEN = "requires_2fa_token" +REQUIRES_2FA_USER = "requires_2fa_user" +VALID_USERS = [AUTHENTICATED_USER, REQUIRES_2FA_USER, APPLE_ID_EMAIL, ICLOUD_ID_EMAIL] VALID_PASSWORD = "valid_password" +VALID_COOKIE = "valid_cookie" +VALID_TOKEN = "valid_token" +VALID_2FA_CODE = "000000" +VALID_TOKENS = [VALID_TOKEN, REQUIRES_2FA_TOKEN] CLIENT_ID = "client_id" diff --git a/tests/const_login.py b/tests/const_login.py index 25ec7c0..352c2ae 100644 --- a/tests/const_login.py +++ b/tests/const_login.py @@ -16,6 +16,10 @@ A_DS_ID = "123456-12-12345678-1234-1234-1234-123456789012" + PERSON_ID WIDGET_KEY = "widget_key" + PERSON_ID # Data +AUTH_OK = { + "authType": "hsa2" +} + LOGIN_WORKING = { "dsInfo": { "lastName": LAST_NAME, @@ -209,7 +213,7 @@ LOGIN_WORKING = { } # Setup data -LOGIN_2SA = { +LOGIN_2FA = { "dsInfo": { "lastName": LAST_NAME, "iCDPEnabled": False, diff --git a/tests/test_cmdline.py b/tests/test_cmdline.py index 1a80239..42c2509 100644 --- a/tests/test_cmdline.py +++ b/tests/test_cmdline.py @@ -2,7 +2,7 @@ """Cmdline tests.""" from pyicloud import cmdline from . import PyiCloudServiceMock -from .const import AUTHENTICATED_USER, REQUIRES_2SA_USER, VALID_PASSWORD +from .const import AUTHENTICATED_USER, REQUIRES_2FA_USER, VALID_PASSWORD, VALID_2FA_CODE from .const_findmyiphone import FMI_FAMILY_WORKING import os @@ -74,16 +74,16 @@ class TestCmdline(TestCase): @patch("keyring.get_password", return_value=None) @patch("pyicloud.cmdline.input") - def test_username_password_requires_2sa( + def test_username_password_requires_2fa( self, mock_input, mock_get_password ): # pylint: disable=unused-argument """Test username and password commands.""" # Valid connection for the first time - mock_input.return_value = "0" + mock_input.return_value = VALID_2FA_CODE with pytest.raises(SystemExit, match="0"): # fmt: off self.main([ - '--username', REQUIRES_2SA_USER, + '--username', REQUIRES_2FA_USER, '--password', VALID_PASSWORD, '--non-interactive', ])