From d056c9b7080cc43c9f818defdd0a28d149f8cdd0 Mon Sep 17 00:00:00 2001
From: Adam Coddington <me@adamcoddington.net>
Date: Sun, 17 May 2015 23:11:27 -0700
Subject: [PATCH] Updating cookie handling to serialize all cookies following
 the request.

---
 pyicloud/base.py    | 70 ++++++++++++++++++++++++++-------------------
 pyicloud/cmdline.py | 10 +++++++
 2 files changed, 51 insertions(+), 29 deletions(-)

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(