-
Notifications
You must be signed in to change notification settings - Fork 60
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Persistent Cookie Jar added as default option for better session handling.
- Loading branch information
1 parent
24a25ab
commit 9b2f404
Showing
3 changed files
with
130 additions
and
3 deletions.
There are no files selected for viewing
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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,5 +1,5 @@ | ||
# -*- coding: utf-8 -*- | ||
|
||
__version__ = "1.3.6" | ||
__version__ = "1.4.0" | ||
|
||
from vulners.api import Vulners |
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
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
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,116 @@ | ||
# -*- coding: utf-8 -*- | ||
# =============================== | ||
# Persistent cookie jar for Requests by Kir Ermakov <[email protected]> | ||
# 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) |