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.
This commit is contained in:
FiZiX 2015-05-13 13:01:08 -04:00
parent 7f7becb7d8
commit 897f905855

View file

@ -3,6 +3,9 @@ import hashlib
import json import json
import requests import requests
import sys import sys
import pickle
import os
from re import match
from pyicloud.exceptions import PyiCloudFailedLoginException from pyicloud.exceptions import PyiCloudFailedLoginException
from pyicloud.services import ( from pyicloud.services import (
@ -37,6 +40,8 @@ class PyiCloudService(object):
self._base_system_url = '%s/system/version.json' % self._home_endpoint self._base_system_url = '%s/system/version.json' % self._home_endpoint
self._base_webauth_url = '%s/refreshWebAuth' % self._push_endpoint self._base_webauth_url = '%s/refreshWebAuth' % self._push_endpoint
self._cookie_directory = 'cookies'
self.session = requests.Session() self.session = requests.Session()
self.session.verify = False self.session.verify = False
self.session.headers.update({ self.session.headers.update({
@ -86,6 +91,24 @@ class PyiCloudService(object):
""" """
self.refresh_validate() 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 = dict(self.user)
data.update({'id': self.params['id'], 'extended_login': False}) data.update({'id': self.params['id'], 'extended_login': False})
req = self.session.post( req = self.session.post(
@ -98,6 +121,14 @@ class PyiCloudService(object):
msg = 'Invalid email/password combination.' msg = 'Invalid email/password combination.'
raise PyiCloudFailedLoginException(msg) 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.refresh_validate()
self.discovery = req.json() self.discovery = req.json()