simono41/pyicloud
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Added tests

Browse source
This commit is contained in:
Niccolò Zapponi 2020-10-29 16:51:08 +00:00
parent b3aee79dcb
commit 4adbfb32ec
No known key found for this signature in database
GPG key ID: 328B304DC670A51E
6 changed files with 104 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

61
pyicloud/base.py
View file

@ -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:
@ -345,14 +344,16 @@ class PyiCloudService(object):
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

17
pyicloud/cmdline.py
View file

@ -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.",

55
tests/__init__.py
View file

@ -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)

9
tests/const.py
View file

@ -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"

6
tests/const_login.py
View file

@ -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,

8
tests/test_cmdline.py
View file

@ -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',
])