From 733071a0f73cdb995c257a48631f23b516e99ee4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tor=20Arne=20Vestb=C3=B8?= <torarnv@gmail.com>
Date: Thu, 7 Jan 2016 20:26:17 +0100
Subject: [PATCH] Rework cookie persistence handling

Instead of trying to look for a specific cookie to save, discarding all
others, and persisting the cookies manually using pickle, we build on
the functionality of cookielib, which already has functionality to save
and load cookiejars. The request library is documented to work with any
subclass of cookielib.CookieJar.

This ensures that we only save persistent cookies (which includes the
X-APPLE-WEB-KB cookie), and skip session cookies, which should make the
code more future proof in case Apple adds more persistent cookies.

This also fixes #44, which was still occurring because we were
persisting the cookies of the request, not the session, and when
logging in with a persisted cookie the resulting request did not
have the X-APPLE-WEB-KB set, so we ended up overwriting the cookie
file with one that didn't contain any X-APPLE-WEB-KB cookie anymore.
---
 pyicloud/base.py | 63 ++++++++++++------------------------------------
 1 file changed, 15 insertions(+), 48 deletions(-)

diff --git a/pyicloud/base.py b/pyicloud/base.py
index a774c07c..0b6faa05 100644
--- a/pyicloud/base.py
+++ b/pyicloud/base.py
@@ -1,4 +1,5 @@
 import copy
+import cookielib
 import uuid
 import hashlib
 import json
@@ -65,6 +66,15 @@ def __init__(self, apple_id, password, cookie_directory=None, verify=True):
             'User-Agent': 'Opera/9.52 (X11; Linux i686; U; en)'
         })
 
+        cookiejar_path = self._get_cookiejar_path()
+        self.session.cookies = cookielib.LWPCookieJar(filename=cookiejar_path)
+        if os.path.exists(cookiejar_path):
+            try:
+                self.session.cookies.load()
+            except cookielib.LoadError:
+                # Most likely a pickled cookiejar from earlier versions
+                pass
+
         self.params = {}
 
         self.authenticate()
@@ -105,15 +115,6 @@ def authenticate(self):
         """
         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(
@@ -126,56 +127,22 @@ def authenticate(self):
             msg = 'Invalid email/password combination.'
             raise PyiCloudFailedLoginException(msg)
 
-        self._update_cookie(req)
+        if not os.path.exists(self._cookie_directory):
+            os.mkdir(self._cookie_directory)
+        self.session.cookies.save()
 
         self.refresh_validate()
 
         self.discovery = req.json()
         self.webservices = self.discovery['webservices']
 
-    def _get_cookie_path(self):
-        # Set path for cookie file
+    def _get_cookiejar_path(self):
+        # Get path for cookiejar 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()
-
-        # We really only want to keep the cookies having names
-        # starting with 'X-APPLE-WEB-KB'
-        for cookie_name, value in request.cookies.items():
-            if not cookie_name.startswith('X-APPLE-WEB-KB'):
-                del request.cookies[cookie_name]
-
-        # Save the cookie in a pickle file
-        with open(cookiefile, 'wb') as f:
-            pickle.dump(request.cookies, f)
-
-        self._cookies = copy.deepcopy(request.cookies)
-
     @property
     def devices(self):
         """ Return all devices."""