From 897f9058554678008fc751e71b24e082b61fdb12 Mon Sep 17 00:00:00 2001
From: FiZiX <FiZiX@users.noreply.github.com>
Date: Wed, 13 May 2015 13:01:08 -0400
Subject: [PATCH] Add X-APPLE-WEB-KB cookie handling

Add X-APPLE-WEB-KB cookie handling to prevent e-mail from Apple every time script is run.
---
 pyicloud/base.py | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/pyicloud/base.py b/pyicloud/base.py
index 739e7bc..bc83cc7 100644
--- a/pyicloud/base.py
+++ b/pyicloud/base.py
@@ -3,6 +3,9 @@ import hashlib
 import json
 import requests
 import sys
+import pickle
+import os
+from re import match
 
 from pyicloud.exceptions import PyiCloudFailedLoginException
 from pyicloud.services import (
@@ -37,6 +40,8 @@ class PyiCloudService(object):
         self._base_system_url = '%s/system/version.json' % self._home_endpoint
         self._base_webauth_url = '%s/refreshWebAuth' % self._push_endpoint
 
+        self._cookie_directory = 'cookies'
+
         self.session = requests.Session()
         self.session.verify = False
         self.session.headers.update({
@@ -86,6 +91,24 @@ 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)
+
+        # Set path for cookie file
+        cookiefile = self.user.get('apple_id')
+        cookiefile = os.path.join(self._cookie_directory, ''.join([c for c in cookiefile if match(r'\w', c)]))
+
+        # Check if cookie file already exists
+        if os.path.isfile(cookiefile):
+            # Get cookie data from file
+            with open(cookiefile, 'rb') as f:
+                webKBCookie = pickle.load(f)
+            self.session.cookies = requests.utils.cookiejar_from_dict(webKBCookie)
+        else:
+            webKBCookie = None
+
         data = dict(self.user)
         data.update({'id': self.params['id'], 'extended_login': False})
         req = self.session.post(
@@ -98,6 +121,14 @@ 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.refresh_validate()
 
         self.discovery = req.json()