Instantly share code, notes, and snippets.
Last active
January 24, 2019 18:41
-
Star
(0)
0
You must be signed in to star a gist -
Fork
(0)
0
You must be signed in to fork a gist
-
Save vulnersCom/565bbfdea21856dcd3a1e73922c6b260 to your computer and use it in GitHub Desktop.
Beta version of persistent cookie module
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# -*- coding: utf-8 -*- | |
# =============================== | |
# Persistent cookie jar for Requests by Kir Ermakov <isox@vulners.com> | |
# It holds cookies in temp file and recovers on process start | |
# | |
# Usage: | |
# | |
# import requests | |
# | |
# opener = requests.session() | |
# opener.cookies = PersistentCookieJar() | |
# opener.get("url") | |
# | |
# | |
# | |
# | |
# =============================== | |
from requests.cookies import RequestsCookieJar | |
import platform | |
import tempfile | |
import json | |
import six | |
import codecs | |
import os | |
import warnings | |
import sys | |
import hashlib | |
class PersistentCookieJar(RequestsCookieJar): | |
""" | |
This Cookie Jar is designed to hold persistent cookies using temp dir. | |
""" | |
def __init__(self, file_path = None, *args, **kwargs): | |
""" | |
Additional parameter - file path location, that can be changed. | |
It will hold session data there. | |
:param file_path: String, expected full path with filename. If NONE is set, it will we be created dynamically. | |
""" | |
super(PersistentCookieJar, self).__init__(*args, **kwargs) | |
self.__file_path = file_path or os.path.join(self.__get_temp_dir(), self.__get_module_name()) | |
# Try to recover from file if it does exist | |
recover_candidate = self.__recover_from_file(self.__file_path) | |
if recover_candidate: | |
self.update(recover_candidate) | |
self.__write_down() | |
def __get_temp_dir(self): | |
""" | |
Internal method for capturing location of the temp path. | |
For MacOS it's hardcoded to use /tmp | |
:return: string, OS tmp path | |
""" | |
return '/tmp' if platform.system() == 'Darwin' else tempfile.gettempdir() | |
def __get_module_name(self): | |
""" | |
Internal method for gathering Python module name. We need it to make some difference for cookie jars created by separate projects. | |
We are taking sys __main__, it's file name and then takes a hash from it. | |
:return: string, Python module name | |
""" | |
full_module_file_path = six.text_type(sys.modules['__main__'].__file__) | |
path_hash = hashlib.sha1(full_module_file_path.encode('utf-8')).hexdigest() | |
return "%s.cookiejar" % path_hash | |
def __write_down(self): | |
""" | |
Internal method for tearing data to disk. | |
:return: None | |
""" | |
with open(self.__file_path, 'wb') as cookie_file: | |
cookie_file.write(codecs.encode(six.text_type(json.dumps(self.get_dict())).encode(), "base64")) | |
def __recover_from_file(self, file_path): | |
""" | |
Recovers self state object from the file. If something fails it will return none and writes down warning | |
:param file_path: State file location | |
:return: recovered PersistentCookieJar object | |
""" | |
if not os.path.exists(file_path): | |
# If it does not exists it's not actually a problem, we will just create a new one | |
return None | |
if not os.path.isfile(file_path): | |
warnings.warn("%s file path %s is not a file" % (self.__class__.__name__, file_path)) | |
return None | |
if not os.access(file_path, os.R_OK): | |
warnings.warn("%s file path %s can not be read" % (self.__class__.__name__, file_path)) | |
return None | |
with open(file_path, "rb") as cookie_file: | |
try: | |
cookie_jar = json.loads(codecs.decode(cookie_file.read(), "base64")) | |
if not isinstance(cookie_jar, dict): | |
warnings.warn("%s recovered object, but it do mismatch with self class. Recovered type is %s" % (self.__class__.__name__, type(cookie_jar))) | |
return None | |
return cookie_jar | |
except Exception as exc: | |
warnings.warn("%s failed to recover session from %s with error %s" % (self.__class__.__name__, file_path, exc)) | |
return None | |
def set_cookie(self, cookie, *args, **kwargs): | |
""" | |
This method is used to find the good moment for tearing down to the disk and save Jar state. | |
On any modification it will save json object state to the disk | |
""" | |
self.__write_down() | |
return super(PersistentCookieJar, self).set_cookie(cookie, *args, **kwargs) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment