Merge pull request #1074 from mitmproxy/move-response-refresh

move HTTPResponse.refresh into netlib

2 files changed, 68 insertions, 0 deletions

diff --git a/netlib/http/cookies.py b/netlib/http/cookies.py
index 18544b5e..caa84ff7 100644
--- a/netlib/http/cookies.py
+++ b/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
diff --git a/netlib/http/response.py b/netlib/http/response.py
index 5fce4d79..efd7f60a 100644
--- a/netlib/http/response.py
+++ b/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