move HTTPResponse.refresh into netlib

mitmproxy/models/http.py

@@ -1,5 +1,4 @@
 from __future__ import (absolute_import, print_function, division)
-from six.moves import http_cookies as Cookie
 import cgi
 import copy
 import warnings
@@ -262,66 +261,6 @@ class HTTPResponse(MessageMixin, Response):
         )
         return resp
 
-    def _refresh_cookie(self, c, delta):
-        """
-            Takes a cookie string c and a time delta in seconds, and returns
-            a refreshed cookie string.
-        """
-        try:
-            c = Cookie.SimpleCookie(str(c))
-        except Cookie.CookieError:
-            raise ValueError("Invalid Cookie")
-        for i in c.values():
-            if "expires" in i:
-                d = parsedate_tz(i["expires"])
-                if d:
-                    d = mktime_tz(d) + delta
-                    i["expires"] = formatdate(d)
-                else:
-                    # This can happen when the expires tag is invalid.
-                    # reddit.com sends a an expires tag like this: "Thu, 31 Dec
-                    # 2037 23:59:59 GMT", which is valid RFC 1123, but not
-                    # strictly correct according to the cookie spec. Browsers
-                    # appear to parse this tolerantly - maybe we should too.
-                    # For now, we just ignore this.
-                    del i["expires"]
-        ret = c.output(header="").strip()
-        if not ret:
-            raise ValueError("Invalid Cookie")
-        return ret
-
-    def refresh(self, now=None):
-        """
-            This fairly complex and heuristic function refreshes a server
-            response for replay.
-
-                - It adjusts date, expires and last-modified headers.
-                - It adjusts cookie expiration.
-        """
-        if not now:
-            now = time.time()
-        delta = now - self.timestamp_start
-        refresh_headers = [
-            "date",
-            "expires",
-            "last-modified",
-        ]
-        for i in refresh_headers:
-            if i in self.headers:
-                d = parsedate_tz(self.headers[i])
-                if d:
-                    new = mktime_tz(d) + delta
-                    self.headers[i] = formatdate(new)
-        c = []
-        for set_cookie_header in self.headers.get_all("set-cookie"):
-            try:
-                refreshed = self._refresh_cookie(set_cookie_header, delta)
-            except ValueError:
-                refreshed = set_cookie_header
-            c.append(refreshed)
-        if c:
-            self.headers.set_all("set-cookie", c)
-
 
 class HTTPFlow(Flow):
 

netlib/http/cookies.py

@@ -1,4 +1,6 @@
+from six.moves import http_cookies as Cookie
 import re
+from email.utils import parsedate_tz, formatdate, mktime_tz
 
 from .. import odict
 
@@ -191,3 +193,35 @@ def format_cookie_header(od):
         Formats a Cookie header value.
     """
     return _format_pairs(od.lst)
+
+
+def refresh_set_cookie_header(c, delta):
+    """
+    Args:
+        c: A Set-Cookie string
+        delta: Time delta in seconds
+    Returns:
+        A refreshed Set-Cookie string
+    """
+    try:
+        c = Cookie.SimpleCookie(str(c))
+    except Cookie.CookieError:
+        raise ValueError("Invalid Cookie")
+    for i in c.values():
+        if "expires" in i:
+            d = parsedate_tz(i["expires"])
+            if d:
+                d = mktime_tz(d) + delta
+                i["expires"] = formatdate(d)
+            else:
+                # This can happen when the expires tag is invalid.
+                # reddit.com sends a an expires tag like this: "Thu, 31 Dec
+                # 2037 23:59:59 GMT", which is valid RFC 1123, but not
+                # strictly correct according to the cookie spec. Browsers
+                # appear to parse this tolerantly - maybe we should too.
+                # For now, we just ignore this.
+                del i["expires"]
+    ret = c.output(header="").strip()
+    if not ret:
+        raise ValueError("Invalid Cookie")
+    return ret

netlib/http/response.py

@@ -1,6 +1,8 @@
 from __future__ import absolute_import, print_function, division
 
 import warnings
+from email.utils import parsedate_tz, formatdate, mktime_tz
+import time
 
 from . import cookies
 from .headers import Headers
@@ -94,6 +96,38 @@ class Response(Message):
             values.append(header)
         self.headers.set_all("set-cookie", values)
 
+    def refresh(self, now=None):
+        """
+        This fairly complex and heuristic function refreshes a server
+        response for replay.
+
+            - It adjusts date, expires and last-modified headers.
+            - It adjusts cookie expiration.
+        """
+        if not now:
+            now = time.time()
+        delta = now - self.timestamp_start
+        refresh_headers = [
+            "date",
+            "expires",
+            "last-modified",
+        ]
+        for i in refresh_headers:
+            if i in self.headers:
+                d = parsedate_tz(self.headers[i])
+                if d:
+                    new = mktime_tz(d) + delta
+                    self.headers[i] = formatdate(new)
+        c = []
+        for set_cookie_header in self.headers.get_all("set-cookie"):
+            try:
+                refreshed = cookies.refresh_set_cookie_header(set_cookie_header, delta)
+            except ValueError:
+                refreshed = set_cookie_header
+            c.append(refreshed)
+        if c:
+            self.headers.set_all("set-cookie", c)
+
     # Legacy
 
     def get_cookies(self):  # pragma: no cover

test/mitmproxy/test_flow.py

@@ -1151,38 +1151,6 @@ class TestResponse:
         resp2 = resp.copy()
         assert resp2.get_state() == resp.get_state()
 
-    def test_refresh(self):
-        r = HTTPResponse.wrap(netlib.tutils.tresp())
-        n = time.time()
-        r.headers["date"] = email.utils.formatdate(n)
-        pre = r.headers["date"]
-        r.refresh(n)
-        assert pre == r.headers["date"]
-        r.refresh(n + 60)
-
-        d = email.utils.parsedate_tz(r.headers["date"])
-        d = email.utils.mktime_tz(d)
-        # Weird that this is not exact...
-        assert abs(60 - (d - n)) <= 1
-
-        r.headers["set-cookie"] = "MOO=BAR; Expires=Tue, 08-Mar-2011 00:20:38 GMT; Path=foo.com; Secure"
-        r.refresh()
-
-    def test_refresh_cookie(self):
-        r = HTTPResponse.wrap(netlib.tutils.tresp())
-
-        # Invalid expires format, sent to us by Reddit.
-        c = "rfoo=bar; Domain=reddit.com; expires=Thu, 31 Dec 2037 23:59:59 GMT; Path=/"
-        assert r._refresh_cookie(c, 60)
-
-        c = "MOO=BAR; Expires=Tue, 08-Mar-2011 00:20:38 GMT; Path=foo.com; Secure"
-        assert "00:21:38" in r._refresh_cookie(c, 60)
-
-        # https://github.com/mitmproxy/mitmproxy/issues/773
-        c = ">=A"
-        with tutils.raises(ValueError):
-            r._refresh_cookie(c, 60)
-
     def test_replace(self):
         r = HTTPResponse.wrap(netlib.tutils.tresp())
         r.headers["Foo"] = "fOo"

test/netlib/http/test_cookies.py

@@ -1,4 +1,5 @@
 from netlib.http import cookies
+from netlib.tutils import raises
 
 
 def test_read_token():
@@ -216,3 +217,18 @@ def test_parse_set_cookie_header():
             assert ret2[2].lst == expected[2]
         else:
             assert ret is None
+
+
+def test_refresh_cookie():
+
+    # Invalid expires format, sent to us by Reddit.
+    c = "rfoo=bar; Domain=reddit.com; expires=Thu, 31 Dec 2037 23:59:59 GMT; Path=/"
+    assert cookies.refresh_set_cookie_header(c, 60)
+
+    c = "MOO=BAR; Expires=Tue, 08-Mar-2011 00:20:38 GMT; Path=foo.com; Secure"
+    assert "00:21:38" in cookies.refresh_set_cookie_header(c, 60)
+
+    # https://github.com/mitmproxy/mitmproxy/issues/773
+    c = ">=A"
+    with raises(ValueError):
+        cookies.refresh_set_cookie_header(c, 60)

test/netlib/http/test_response.py

@@ -1,6 +1,9 @@
 from __future__ import absolute_import, print_function, division
 
+import email
+
 import six
+import time
 
 from netlib.http import Headers
 from netlib.odict import ODict, ODictCaseless
@@ -100,3 +103,23 @@ class TestResponseUtils(object):
         v = resp.cookies
         assert len(v) == 1
         assert v["foo"] == [["bar", ODictCaseless()]]
+
+    def test_refresh(self):
+        r = tresp()
+        n = time.time()
+        r.headers["date"] = email.utils.formatdate(n)
+        pre = r.headers["date"]
+        r.refresh(n)
+        assert pre == r.headers["date"]
+        r.refresh(n + 60)
+
+        d = email.utils.parsedate_tz(r.headers["date"])
+        d = email.utils.mktime_tz(d)
+        # Weird that this is not exact...
+        assert abs(60 - (d - n)) <= 1
+
+        cookie = "MOO=BAR; Expires=Tue, 08-Mar-2011 00:20:38 GMT; Path=foo.com; Secure"
+        r.headers["set-cookie"] = cookie
+        r.refresh()
+        # Cookie refreshing is tested in test_cookies, we just make sure that it's triggered here.
+        assert cookie != r.headers["set-cookie"]