141 lines
4.5 KiB
Python
141 lines
4.5 KiB
Python
import re
|
|
import string
|
|
|
|
# ------------------------------------------------------------ #
|
|
# SimpleCookie
|
|
# ------------------------------------------------------------ #
|
|
|
|
# Straight up copied this section of dark magic from SimpleCookie
|
|
|
|
_LegalChars = string.ascii_letters + string.digits + "!#$%&'*+-.^_`|~:"
|
|
_UnescapedChars = _LegalChars + ' ()/<=>?@[]{}'
|
|
|
|
_Translator = {n: '\\%03o' % n
|
|
for n in set(range(256)) - set(map(ord, _UnescapedChars))}
|
|
_Translator.update({
|
|
ord('"'): '\\"',
|
|
ord('\\'): '\\\\',
|
|
})
|
|
|
|
|
|
def _quote(str):
|
|
"""Quote a string for use in a cookie header.
|
|
If the string does not need to be double-quoted, then just return the
|
|
string. Otherwise, surround the string in doublequotes and quote
|
|
(with a \) special characters.
|
|
"""
|
|
if str is None or _is_legal_key(str):
|
|
return str
|
|
else:
|
|
return '"' + str.translate(_Translator) + '"'
|
|
|
|
|
|
_is_legal_key = re.compile('[%s]+' % re.escape(_LegalChars)).fullmatch
|
|
|
|
# ------------------------------------------------------------ #
|
|
# Custom SimpleCookie
|
|
# ------------------------------------------------------------ #
|
|
|
|
|
|
class CookieJar(dict):
|
|
"""CookieJar dynamically writes headers as cookies are added and removed
|
|
It gets around the limitation of one header per name by using the
|
|
MultiHeader class to provide a unique key that encodes to Set-Cookie.
|
|
"""
|
|
|
|
def __init__(self, headers):
|
|
super().__init__()
|
|
self.headers = headers
|
|
self.cookie_headers = {}
|
|
|
|
def __setitem__(self, key, value):
|
|
# If this cookie doesn't exist, add it to the header keys
|
|
cookie_header = self.cookie_headers.get(key)
|
|
if not cookie_header:
|
|
cookie = Cookie(key, value)
|
|
cookie['path'] = '/'
|
|
cookie_header = MultiHeader("Set-Cookie")
|
|
self.cookie_headers[key] = cookie_header
|
|
self.headers[cookie_header] = cookie
|
|
return super().__setitem__(key, cookie)
|
|
else:
|
|
self[key].value = value
|
|
|
|
def __delitem__(self, key):
|
|
if key not in self.cookie_headers:
|
|
self[key] = ''
|
|
self[key]['max-age'] = 0
|
|
else:
|
|
cookie_header = self.cookie_headers[key]
|
|
del self.headers[cookie_header]
|
|
del self.cookie_headers[key]
|
|
return super().__delitem__(key)
|
|
|
|
|
|
class Cookie(dict):
|
|
"""A stripped down version of Morsel from SimpleCookie #gottagofast"""
|
|
_keys = {
|
|
"expires": "expires",
|
|
"path": "Path",
|
|
"comment": "Comment",
|
|
"domain": "Domain",
|
|
"max-age": "Max-Age",
|
|
"secure": "Secure",
|
|
"httponly": "HttpOnly",
|
|
"version": "Version",
|
|
}
|
|
_flags = {'secure', 'httponly'}
|
|
|
|
def __init__(self, key, value):
|
|
if key in self._keys:
|
|
raise KeyError("Cookie name is a reserved word")
|
|
if not _is_legal_key(key):
|
|
raise KeyError("Cookie key contains illegal characters")
|
|
self.key = key
|
|
self.value = value
|
|
super().__init__()
|
|
|
|
def __setitem__(self, key, value):
|
|
if key not in self._keys:
|
|
raise KeyError("Unknown cookie property")
|
|
return super().__setitem__(key, value)
|
|
|
|
def encode(self, encoding):
|
|
output = ['%s=%s' % (self.key, _quote(self.value))]
|
|
for key, value in self.items():
|
|
if key == 'max-age':
|
|
try:
|
|
output.append('%s=%d' % (self._keys[key], value))
|
|
except TypeError:
|
|
output.append('%s=%s' % (self._keys[key], value))
|
|
elif key == 'expires':
|
|
try:
|
|
output.append('%s=%s' % (
|
|
self._keys[key],
|
|
value.strftime("%a, %d-%b-%Y %T GMT")
|
|
))
|
|
except AttributeError:
|
|
output.append('%s=%s' % (self._keys[key], value))
|
|
elif key in self._flags:
|
|
if self[key]:
|
|
output.append(self._keys[key])
|
|
else:
|
|
output.append('%s=%s' % (self._keys[key], value))
|
|
|
|
return "; ".join(output).encode(encoding)
|
|
|
|
# ------------------------------------------------------------ #
|
|
# Header Trickery
|
|
# ------------------------------------------------------------ #
|
|
|
|
|
|
class MultiHeader:
|
|
"""String-holding object which allow us to set a header within response
|
|
that has a unique key, but may contain duplicate header names
|
|
"""
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
def encode(self):
|
|
return self.name.encode()
|