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.
This commit is contained in:
Tor Arne Vestbø 2016-01-07 20:26:17 +01:00
parent d2bbf7dee5
commit 733071a0f7

View file

@ -1,4 +1,5 @@
import copy import copy
import cookielib
import uuid import uuid
import hashlib import hashlib
import json import json
@ -65,6 +66,15 @@ class PyiCloudService(object):
'User-Agent': 'Opera/9.52 (X11; Linux i686; U; en)' '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.params = {}
self.authenticate() self.authenticate()
@ -105,15 +115,6 @@ 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)
cookie = self._get_cookie()
if cookie:
self.session.cookies = cookie
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(
@ -126,56 +127,22 @@ class PyiCloudService(object):
msg = 'Invalid email/password combination.' msg = 'Invalid email/password combination.'
raise PyiCloudFailedLoginException(msg) 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.refresh_validate()
self.discovery = req.json() self.discovery = req.json()
self.webservices = self.discovery['webservices'] self.webservices = self.discovery['webservices']
def _get_cookie_path(self): def _get_cookiejar_path(self):
# Set path for cookie file # Get path for cookiejar file
return os.path.join( return os.path.join(
self._cookie_directory, self._cookie_directory,
''.join([c for c in self.user.get('apple_id') if match(r'\w', c)]) ''.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 @property
def devices(self): def devices(self):
""" Return all devices.""" """ Return all devices."""