fix #2250, add type info to cookie module

This commit is contained in:
Maximilian Hils 2017-04-26 14:15:33 +02:00
parent 7607240c30
commit 87610cc8b2
4 changed files with 57 additions and 49 deletions

View File

@ -206,7 +206,7 @@ def format_request_cookies(fields):
def format_response_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): def name_value(obj):

View File

@ -1,7 +1,7 @@
import collections
import email.utils import email.utils
import re import re
import time import time
from typing import Tuple, List, Iterable
from mitmproxy.types import multidict 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 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"([\"\\])") ESCAPE = re.compile(r"([\"\\])")
@ -43,7 +40,8 @@ class CookieAttrs(multidict.MultiDict):
return values[-1] 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): def _read_until(s, start, term):
@ -131,15 +129,15 @@ def _read_cookie_pairs(s, off=0):
return pairs, off 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. Read pairs of lhs=rhs values from SetCookie headers while handling multiple cookies.
off: start offset off: start offset
specials: attributes that are treated specially specials: attributes that are treated specially
""" """
cookies = [] cookies = [] # type: List[TPairs]
pairs = [] pairs = [] # type: TPairs
while True: while True:
lhs, off = _read_key(s, off, ";=,") lhs, off = _read_key(s, off, ";=,")
@ -182,7 +180,7 @@ def _read_set_cookie_pairs(s, off=0):
return cookies, off return cookies, off
def _has_special(s): def _has_special(s: str) -> bool:
for i in s: for i in s:
if i in '",;\\': if i in '",;\\':
return True return True
@ -238,41 +236,44 @@ def format_cookie_header(lst):
return _format_pairs(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 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 CookieAttrs dict of attributes. No attempt is made to parse attribute
values - they are treated purely as strings. values - they are treated purely as strings.
""" """
cookie_pairs, off = _read_set_cookie_pairs(line) 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 return cookies
def parse_set_cookie_headers(headers): def parse_set_cookie_headers(headers: Iterable[str]) -> List[TSetCookie]:
rv = [] rv = []
for header in headers: for header in headers:
cookies = parse_set_cookie_header(header) 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 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. Formats a Set-Cookie header value.
""" """
rv = [] rv = []
for set_cookie in set_cookies: for name, value, attrs in set_cookies:
name, value, attrs = set_cookie
pairs = [(name, value)] pairs = [(name, value)]
pairs.extend( pairs.extend(
@ -284,16 +285,19 @@ def format_set_cookie_header(set_cookies):
return ", ".join(rv) return ", ".join(rv)
def refresh_set_cookie_header(c, delta): def refresh_set_cookie_header(c: str, delta: int) -> str:
""" """
Args: Args:
c: A Set-Cookie string c: A Set-Cookie string
delta: Time delta in seconds delta: Time delta in seconds
Returns: Returns:
A refreshed Set-Cookie string 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: if not name or not value:
raise ValueError("Invalid Cookie") raise ValueError("Invalid Cookie")
@ -310,11 +314,7 @@ def refresh_set_cookie_header(c, delta):
# appear to parse this tolerantly - maybe we should too. # appear to parse this tolerantly - maybe we should too.
# For now, we just ignore this. # For now, we just ignore this.
del attrs["expires"] 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): def get_expiration_ts(cookie_attrs):

View File

@ -131,7 +131,11 @@ class Response(message.Message):
def _get_cookies(self): def _get_cookies(self):
h = self.headers.get_all("set-cookie") 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): def _set_cookies(self, value):
cookie_headers = [] cookie_headers = []

View File

@ -283,6 +283,10 @@ def test_refresh_cookie():
c = "foo/bar=bla" c = "foo/bar=bla"
assert cookies.refresh_set_cookie_header(c, 0) 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') @mock.patch('time.time')
def test_get_expiration_ts(*args): def test_get_expiration_ts(*args):