Merge pull request #2269 from mhils/issue-2250

fix #2250, add type info to cookie module

examples/complex/har_dump.py

@@ -206,7 +206,7 @@ def format_request_cookies(fields):
 
 
 def format_response_cookies(fields):
-    return format_cookies((c[0], c[1].value, c[1].attrs) for c in fields)
+    return format_cookies((c[0], c[1][0], c[1][1]) for c in fields)
 
 
 def name_value(obj):

mitmproxy/net/http/cookies.py

@@ -1,7 +1,7 @@
-import collections
 import email.utils
 import re
 import time
+from typing import Tuple, List, Iterable
 
 from mitmproxy.types import multidict
 
@@ -23,10 +23,7 @@ cookies to be set in a single header. Serialization follows RFC6265.
     http://tools.ietf.org/html/rfc2965
 """
 
-_cookie_params = set((
+_cookie_params = {'expires', 'path', 'comment', 'max-age', 'secure', 'httponly', 'version'}
-    'expires', 'path', 'comment', 'max-age',
-    'secure', 'httponly', 'version',
-))
 
 ESCAPE = re.compile(r"([\"\\])")
 
@@ -43,7 +40,8 @@ class CookieAttrs(multidict.MultiDict):
         return values[-1]
 
 
-SetCookie = collections.namedtuple("SetCookie", ["value", "attrs"])
+TSetCookie = Tuple[str, str, CookieAttrs]
+TPairs = List[List[str]]  # TODO: Should be List[Tuple[str,str]]?
 
 
 def _read_until(s, start, term):
@@ -131,15 +129,15 @@ def _read_cookie_pairs(s, off=0):
     return pairs, off
 
 
-def _read_set_cookie_pairs(s, off=0):
+def _read_set_cookie_pairs(s: str, off=0) -> Tuple[List[TPairs], int]:
     """
         Read pairs of lhs=rhs values from SetCookie headers while handling multiple cookies.
 
         off: start offset
         specials: attributes that are treated specially
     """
-    cookies = []
+    cookies = []  # type: List[TPairs]
-    pairs = []
+    pairs = []  # type: TPairs
 
     while True:
         lhs, off = _read_key(s, off, ";=,")
@@ -182,7 +180,7 @@ def _read_set_cookie_pairs(s, off=0):
     return cookies, off
 
 
-def _has_special(s):
+def _has_special(s: str) -> bool:
     for i in s:
         if i in '",;\\':
             return True
@@ -238,41 +236,44 @@ def format_cookie_header(lst):
     return _format_pairs(lst)
 
 
-def parse_set_cookie_header(line):
+def parse_set_cookie_header(line: str) -> List[TSetCookie]:
     """
     Parse a Set-Cookie header value
 
-        Returns a list of (name, value, attrs) tuples, where attrs is a
+    Returns:
+        A list of (name, value, attrs) tuples, where attrs is a
         CookieAttrs dict of attributes. No attempt is made to parse attribute
         values - they are treated purely as strings.
     """
     cookie_pairs, off = _read_set_cookie_pairs(line)
-    cookies = [
+    cookies = []
-        (pairs[0][0], pairs[0][1], CookieAttrs(tuple(x) for x in pairs[1:]))
+    for pairs in cookie_pairs:
-        for pairs in cookie_pairs if pairs
+        if pairs:
-    ]
+            cookie, *attrs = pairs
+            cookies.append((
+                cookie[0],
+                cookie[1],
+                CookieAttrs(attrs)
+            ))
     return cookies
 
 
-def parse_set_cookie_headers(headers):
+def parse_set_cookie_headers(headers: Iterable[str]) -> List[TSetCookie]:
     rv = []
     for header in headers:
         cookies = parse_set_cookie_header(header)
-        if cookies:
+        rv.extend(cookies)
-            for name, value, attrs in cookies:
-                rv.append((name, SetCookie(value, attrs)))
     return rv
 
 
-def format_set_cookie_header(set_cookies):
+def format_set_cookie_header(set_cookies: List[TSetCookie]) -> str:
     """
         Formats a Set-Cookie header value.
     """
 
     rv = []
 
-    for set_cookie in set_cookies:
+    for name, value, attrs in set_cookies:
-        name, value, attrs = set_cookie
 
         pairs = [(name, value)]
         pairs.extend(
@@ -284,16 +285,19 @@ def format_set_cookie_header(set_cookies):
     return ", ".join(rv)
 
 
-def refresh_set_cookie_header(c, delta):
+def refresh_set_cookie_header(c: str, delta: int) -> str:
     """
     Args:
         c: A Set-Cookie string
         delta: Time delta in seconds
     Returns:
         A refreshed Set-Cookie string
+    Raises:
+        ValueError, if the cookie is invalid.
     """
-
+    cookies = parse_set_cookie_header(c)
-    name, value, attrs = parse_set_cookie_header(c)[0]
+    for cookie in cookies:
+        name, value, attrs = cookie
         if not name or not value:
             raise ValueError("Invalid Cookie")
 
@@ -310,11 +314,7 @@ def refresh_set_cookie_header(c, delta):
                 # appear to parse this tolerantly - maybe we should too.
                 # For now, we just ignore this.
                 del attrs["expires"]
-
+    return format_set_cookie_header(cookies)
-    rv = format_set_cookie_header([(name, value, attrs)])
-    if not rv:
-        raise ValueError("Invalid Cookie")
-    return rv
 
 
 def get_expiration_ts(cookie_attrs):

mitmproxy/net/http/response.py

@@ -131,7 +131,11 @@ class Response(message.Message):
 
     def _get_cookies(self):
         h = self.headers.get_all("set-cookie")
-        return tuple(cookies.parse_set_cookie_headers(h))
+        all_cookies = cookies.parse_set_cookie_headers(h)
+        return tuple(
+            (name, (value, attrs))
+            for name, value, attrs in all_cookies
+        )
 
     def _set_cookies(self, value):
         cookie_headers = []

test/mitmproxy/net/http/test_cookies.py

@@ -283,6 +283,10 @@ def test_refresh_cookie():
     c = "foo/bar=bla"
     assert cookies.refresh_set_cookie_header(c, 0)
 
+    # https://github.com/mitmproxy/mitmproxy/issues/2250
+    c = ""
+    assert cookies.refresh_set_cookie_header(c, 60) == ""
+
 
 @mock.patch('time.time')
 def test_get_expiration_ts(*args):