diff --git a/pyicloud/base.py b/pyicloud/base.py index 87b3338..ba9eae8 100644 --- a/pyicloud/base.py +++ b/pyicloud/base.py @@ -109,28 +109,9 @@ class PyiCloudService(object): # If not, create it os.mkdir(self._cookie_directory) - # Set path for cookie file - cookiefile = os.path.join( - self._cookie_directory, - ''.join([c for c in self.user.get('apple_id') if match(r'\w', c)]) - ) - - webKBCookie = None - # Check if cookie file already exists - try: - # Get cookie data from file - with open(cookiefile, 'rb') as f: - webKBCookie = pickle.load(f) - self.session.cookies = requests.utils.cookiejar_from_dict( - webKBCookie - ) - except IOError: - # This just means that the file doesn't exist; that's OK! - pass - except Exception as e: - logger.exception( - "Unexpected error occurred while loading cookies: %s" % (e, ) - ) + cookie = self._get_cookie() + if cookie: + self.session.cookies = cookie data = dict(self.user) data.update({'id': self.params['id'], 'extended_login': False}) @@ -144,19 +125,50 @@ class PyiCloudService(object): msg = 'Invalid email/password combination.' raise PyiCloudFailedLoginException(msg) - # Pull X-APPLE-WEB-KB cookie from cookies - newWebKBCookie = next(({key:val} for key, val in req.cookies.items() if 'X-APPLE-WEB-KB' in key), None) - # Check if cookie changed - if newWebKBCookie and newWebKBCookie != webKBCookie: - # Save the cookie in a pickle file - with open(cookiefile, 'wb') as f: - pickle.dump(newWebKBCookie, f) + self._update_cookie(req) self.refresh_validate() self.discovery = req.json() self.webservices = self.discovery['webservices'] + def _get_cookie_path(self): + # Set path for cookie file + return os.path.join( + self._cookie_directory, + ''.join([c for c in self.user.get('apple_id') if match(r'\w', c)]) + ) + + def _get_cookie(self): + if hasattr(self, '_cookies'): + return self._cookies + + cookiefile = self._get_cookie_path() + + # Check if cookie file already exists + try: + # Get cookie data from file + with open(cookiefile, 'rb') as f: + return pickle.load(f) + except IOError: + # This just means that the file doesn't exist; that's OK! + pass + except Exception as e: + logger.exception( + "Unexpected error occurred while loading cookies: %s" % (e, ) + ) + + return None + + def _update_cookie(self, request): + cookiefile = self._get_cookie_path() + + # Save the cookie in a pickle file + with open(cookiefile, 'wb') as f: + pickle.dump(request.cookies, f) + + self._cookies = request.cookies + @property def devices(self): """ Return all devices.""" diff --git a/pyicloud/cmdline.py b/pyicloud/cmdline.py index e7a1f70..baff280 100644 --- a/pyicloud/cmdline.py +++ b/pyicloud/cmdline.py @@ -6,6 +6,7 @@ command line scripts, and related. """ from __future__ import print_function import argparse +import logging import pickle import sys @@ -150,11 +151,20 @@ def main(args=None): default="", help="Save device data to a file in the current directory.", ) + parser.add_argument( + '--loglevel', + default='INFO', + help='Increase logging verbosity by specifying DEBUG' + ) command_line = parser.parse_args(args) if not command_line.username or not command_line.password: parser.error('No username or password supplied') + logging.basicConfig( + level=logging.getLevelName(command_line.loglevel) + ) + from pyicloud import PyiCloudService try: api = PyiCloudService(