codekingpro/portable-devtools

1import email.utils2import re3import time4from collections.abc import Iterable5 6from mitmproxy.coretypes import multidict7 8"""9A flexible module for cookie parsing and manipulation.10 11This module differs from usual standards-compliant cookie modules in a number12of ways. We try to be as permissive as possible, and to retain even mal-formed13information. Duplicate cookies are preserved in parsing, and can be set in14formatting. We do attempt to escape and quote values where needed, but will not15reject data that violate the specs.16 17Parsing accepts the formats in RFC6265 and partially RFC2109 and RFC2965. We18also parse the comma-separated variant of Set-Cookie that allows multiple19cookies to be set in a single header. Serialization follows RFC6265.20 21    http://tools.ietf.org/html/rfc626522    http://tools.ietf.org/html/rfc210923    http://tools.ietf.org/html/rfc296524"""25 26_cookie_params = {27    "expires",28    "path",29    "comment",30    "max-age",31    "secure",32    "httponly",33    "version",34}35 36ESCAPE = re.compile(r"([\"\\])")37 38 39class CookieAttrs(multidict.MultiDict):40    @staticmethod41    def _kconv(key):42        return key.lower()43 44    @staticmethod45    def _reduce_values(values):46        # See the StickyCookieTest for a weird cookie that only makes sense47        # if we take the last part.48        return values[-1]49 50 51TSetCookie = tuple[str, str | None, CookieAttrs]52TPairs = list[tuple[str, str | None]]53 54 55def _read_until(s, start, term):56    """57    Read until one of the characters in term is reached.58    """59    if start == len(s):60        return "", start + 161    for i in range(start, len(s)):62        if s[i] in term:63            return s[start:i], i64    return s[start : i + 1], i + 165 66 67def _read_quoted_string(s, start):68    """69    start: offset to the first quote of the string to be read70 71    A sort of loose super-set of the various quoted string specifications.72 73    RFC6265 disallows backslashes or double quotes within quoted strings.74    Prior RFCs use backslashes to escape. This leaves us free to apply75    backslash escaping by default and be compatible with everything.76    """77    escaping = False78    ret = []79    # Skip the first quote80    i = start  # initialize in case the loop doesn't run.81    for i in range(start + 1, len(s)):82        if escaping:83            ret.append(s[i])84            escaping = False85        elif s[i] == '"':86            break87        elif s[i] == "\\":88            escaping = True89        else:90            ret.append(s[i])91    return "".join(ret), i + 192 93 94def _read_key(s, start, delims=";="):95    """96    Read a key - the LHS of a token/value pair in a cookie.97    """98    return _read_until(s, start, delims)99 100 101def _read_value(s, start, delims):102    """103    Reads a value - the RHS of a token/value pair in a cookie.104    """105    if start >= len(s):106        return "", start107    elif s[start] == '"':108        return _read_quoted_string(s, start)109    else:110        return _read_until(s, start, delims)111 112 113def _read_cookie_pairs(s, off=0):114    """115    Read pairs of lhs=rhs values from Cookie headers.116 117    off: start offset118    """119    pairs = []120 121    while True:122        lhs, off = _read_key(s, off)123        lhs = lhs.lstrip()124 125        rhs = ""126        if off < len(s) and s[off] == "=":127            rhs, off = _read_value(s, off + 1, ";")128        if rhs or lhs:129            pairs.append([lhs, rhs])130 131        off += 1132 133        if not off < len(s):134            break135 136    return pairs, off137 138 139def _read_set_cookie_pairs(s: str, off=0) -> tuple[list[TPairs], int]:140    """141    Read pairs of lhs=rhs values from SetCookie headers while handling multiple cookies.142 143    off: start offset144    specials: attributes that are treated specially145    """146    cookies: list[TPairs] = []147    pairs: TPairs = []148 149    while True:150        lhs, off = _read_key(s, off, ";=,")151        lhs = lhs.lstrip()152 153        rhs = ""154        if off < len(s) and s[off] == "=":155            rhs, off = _read_value(s, off + 1, ";,")156 157            # Special handling of attributes158            if lhs.lower() == "expires":159                # 'expires' values can contain commas in them so they need to160                # be handled separately.161 162                # We actually bank on the fact that the expires value WILL163                # contain a comma. Things will fail, if they don't.164 165                # '3' is just a heuristic we use to determine whether we've166                # only read a part of the expires value and we should read more.167                if len(rhs) <= 3:168                    trail, off = _read_value(s, off + 1, ";,")169                    rhs = rhs + "," + trail170 171            # as long as there's a "=", we consider it a pair172            pairs.append((lhs, rhs))173 174        elif lhs:175            pairs.append((lhs, None))176 177        # comma marks the beginning of a new cookie178        if off < len(s) and s[off] == ",":179            cookies.append(pairs)180            pairs = []181 182        off += 1183 184        if not off < len(s):185            break186 187    if pairs or not cookies:188        cookies.append(pairs)189 190    return cookies, off191 192 193def _has_special(s: str) -> bool:194    for i in s:195        if i in '",;\\':196            return True197        o = ord(i)198        if o < 0x21 or o > 0x7E:199            return True200    return False201 202 203def _format_pairs(pairs, specials=(), sep="; "):204    """205    specials: A lower-cased list of keys that will not be quoted.206    """207    vals = []208    for k, v in pairs:209        if v is None:210            val = k211        elif k.lower() not in specials and _has_special(v):212            v = ESCAPE.sub(r"\\\1", v)213            v = '"%s"' % v214            val = f"{k}={v}"215        else:216            val = f"{k}={v}"217        vals.append(val)218    return sep.join(vals)219 220 221def _format_set_cookie_pairs(lst):222    return _format_pairs(lst, specials=("expires", "path"))223 224 225def parse_cookie_header(line):226    """227    Parse a Cookie header value.228    Returns a list of (lhs, rhs) tuples.229    """230    pairs, off_ = _read_cookie_pairs(line)231    return pairs232 233 234def parse_cookie_headers(cookie_headers):235    cookie_list = []236    for header in cookie_headers:237        cookie_list.extend(parse_cookie_header(header))238    return cookie_list239 240 241def format_cookie_header(lst):242    """243    Formats a Cookie header value.244    """245    return _format_pairs(lst)246 247 248def parse_set_cookie_header(line: str) -> list[TSetCookie]:249    """250    Parse a Set-Cookie header value251 252    Returns:253        A list of (name, value, attrs) tuples, where attrs is a254        CookieAttrs dict of attributes. No attempt is made to parse attribute255        values - they are treated purely as strings.256    """257    cookie_pairs, off = _read_set_cookie_pairs(line)258    cookies = []259    for pairs in cookie_pairs:260        if pairs:261            cookie, *attrs = pairs262            cookies.append((cookie[0], cookie[1], CookieAttrs(attrs)))263    return cookies264 265 266def parse_set_cookie_headers(headers: Iterable[str]) -> list[TSetCookie]:267    rv = []268    for header in headers:269        cookies = parse_set_cookie_header(header)270        rv.extend(cookies)271    return rv272 273 274def format_set_cookie_header(set_cookies: list[TSetCookie]) -> str:275    """276    Formats a Set-Cookie header value.277    """278 279    rv = []280 281    for name, value, attrs in set_cookies:282        pairs = [(name, value)]283        pairs.extend(attrs.fields if hasattr(attrs, "fields") else attrs)284 285        rv.append(_format_set_cookie_pairs(pairs))286 287    return ", ".join(rv)288 289 290def refresh_set_cookie_header(c: str, delta: int) -> str:291    """292    Args:293        c: A Set-Cookie string294        delta: Time delta in seconds295    Returns:296        A refreshed Set-Cookie string297    Raises:298        ValueError, if the cookie is invalid.299    """300    cookies = parse_set_cookie_header(c)301    for cookie in cookies:302        name, value, attrs = cookie303        if not name or not value:304            raise ValueError("Invalid Cookie")305 306        if "expires" in attrs:307            e = email.utils.parsedate_tz(attrs["expires"])308            if e:309                f = email.utils.mktime_tz(e) + delta310                attrs.set_all("expires", [email.utils.formatdate(f, usegmt=True)])311            else:312                # This can happen when the expires tag is invalid.313                # reddit.com sends a an expires tag like this: "Thu, 31 Dec314                # 2037 23:59:59 GMT", which is valid RFC 1123, but not315                # strictly correct according to the cookie spec. Browsers316                # appear to parse this tolerantly - maybe we should too.317                # For now, we just ignore this.318                del attrs["expires"]319    return format_set_cookie_header(cookies)320 321 322def get_expiration_ts(cookie_attrs):323    """324    Determines the time when the cookie will be expired.325 326    Considering both 'expires' and 'max-age' parameters.327 328    Returns: timestamp of when the cookie will expire.329             None, if no expiration time is set.330    """331    if "expires" in cookie_attrs:332        e = email.utils.parsedate_tz(cookie_attrs["expires"])333        if e:334            return email.utils.mktime_tz(e)335 336    elif "max-age" in cookie_attrs:337        try:338            max_age = int(cookie_attrs["Max-Age"])339        except ValueError:340            pass341        else:342            now_ts = time.time()343            return now_ts + max_age344 345    return None346 347 348def is_expired(cookie_attrs):349    """350    Determines whether a cookie has expired.351 352    Returns: boolean353    """354 355    exp_ts = get_expiration_ts(cookie_attrs)356    now_ts = time.time()357 358    # If no expiration information was provided with the cookie359    if exp_ts is None:360        return False361    else:362        return exp_ts <= now_ts363 364 365def group_cookies(pairs):366    """367    Converts a list of pairs to a (name, value, attrs) for each cookie.368    """369 370    if not pairs:371        return []372 373    cookie_list = []374 375    # First pair is always a new cookie376    name, value = pairs[0]377    attrs = []378 379    for k, v in pairs[1:]:380        if k.lower() in _cookie_params:381            attrs.append((k, v))382        else:383            cookie_list.append((name, value, CookieAttrs(attrs)))384            name, value, attrs = k, v, []385 386    cookie_list.append((name, value, CookieAttrs(attrs)))387    return cookie_list388