cookies.py

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):
    r"""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 = {}
        self.header_key = "Set-Cookie"

    def __setitem__(self, key, value):
        # If this cookie doesn't exist, add it to the header keys
        if not self.cookie_headers.get(key):
            cookie = Cookie(key, value)
            cookie["path"] = "/"
            self.cookie_headers[key] = self.header_key
            self.headers.add(self.header_key, 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]
            # remove it from header
            cookies = self.headers.popall(cookie_header)
            for cookie in cookies:
                if cookie.key != key:
                    self.headers.add(cookie_header, cookie)
            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",
        "samesite": "SameSite",
    }
    _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")
        if value is not False:
            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 and self[key]:
                output.append(self._keys[key])
            else:
                output.append("%s=%s" % (self._keys[key], value))

        return "; ".join(output).encode(encoding)