Merge branch 'patch-1'; Adds functionality which should limit the number of e-mail notices from Apple.

pyicloud/base.py

@@ -1,8 +1,13 @@
 import uuid
 import hashlib
 import json
+import logging
+import pickle
 import requests
 import sys
+import tempfile
+import os
+from re import match
 
 from pyicloud.exceptions import PyiCloudFailedLoginException
 from pyicloud.services import (
@@ -13,6 +18,9 @@ from pyicloud.services import (
 )
 
 
+logger = logging.getLogger(__name__)
+
+
 class PyiCloudService(object):
     """
     A base authentication class for the iCloud service. Handles the
@@ -23,7 +31,7 @@ class PyiCloudService(object):
         pyicloud = PyiCloudService('username@apple.com', 'password')
         pyicloud.iphone.location()
     """
-    def __init__(self, apple_id, password):
+    def __init__(self, apple_id, password, cookie_directory=None):
         self.discovery = None
         self.client_id = str(uuid.uuid1()).upper()
         self.user = {'apple_id': apple_id, 'password': password}
@@ -37,6 +45,16 @@ class PyiCloudService(object):
         self._base_system_url = '%s/system/version.json' % self._home_endpoint
         self._base_webauth_url = '%s/refreshWebAuth' % self._push_endpoint
 
+        if cookie_directory:
+            self._cookie_directory = os.path.expanduser(
+                os.path.normpath(cookie_directory)
+            )
+        else:
+            self._cookie_directory = os.path.join(
+                tempfile.gettempdir(),
+                'pyicloud',
+            )
+
         self.session = requests.Session()
         self.session.verify = False
         self.session.headers.update({
@@ -86,6 +104,15 @@ class PyiCloudService(object):
         """
         self.refresh_validate()
 
+        # Check if cookies directory exists
+        if not os.path.exists(self._cookie_directory):
+            # If not, create it
+            os.mkdir(self._cookie_directory)
+
+        cookie = self._get_cookie()
+        if cookie:
+            self.session.cookies = cookie
+
         data = dict(self.user)
         data.update({'id': self.params['id'], 'extended_login': False})
         req = self.session.post(
@@ -98,11 +125,50 @@ class PyiCloudService(object):
             msg = 'Invalid email/password combination.'
             raise PyiCloudFailedLoginException(msg)
 
+        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."""