From ae290587f6f44ab90572e3435f54293cf63aee81 Mon Sep 17 00:00:00 2001
From: Ilya Kreymer <ikreymer@gmail.com>
Date: Fri, 1 Jul 2016 10:15:59 -0400
Subject: [PATCH] temp cookie store: add add_cookie() function for explicitly
 adding cookie, make expiry configurable related to webrecorder/webrecorder#79

---
 urlrewrite/cookies.py | 54 +++++++++++++++++++++++++++----------------
 1 file changed, 34 insertions(+), 20 deletions(-)

diff --git a/urlrewrite/cookies.py b/urlrewrite/cookies.py
index 5ebe1144..9823da47 100644
--- a/urlrewrite/cookies.py
+++ b/urlrewrite/cookies.py
@@ -1,6 +1,7 @@
 from pywb.rewrite.cookie_rewriter import WbUrlBaseCookieRewriter
 from pywb.utils.timeutils import datetime_to_http_date
 from six.moves.http_cookiejar import CookieJar, DefaultCookiePolicy
+from six.moves import zip
 
 import redis
 
@@ -12,42 +13,54 @@ import six
 
 # =============================================================================
 class CookieTracker(object):
-    def __init__(self, redis):
+    def __init__(self, redis, expire_time=120):
         self.redis = redis
+        self.expire_time = expire_time
 
     def get_rewriter(self, url_rewriter, cookie_key):
-        return DomainCacheCookieRewriter(url_rewriter,
-                                         self.redis,
-                                         cookie_key)
+        return DomainCacheCookieRewriter(url_rewriter, self, cookie_key)
 
     def get_cookie_headers(self, url, cookie_key):
         subds = self.get_subdomains(url)
+
         if not subds:
             return None, None
 
         with redis.utils.pipeline(self.redis) as pi:
-            for x in subds:
-                pi.hgetall(cookie_key + '.' + x)
+            for domain in subds:
+                pi.hgetall(cookie_key + '.' + domain)
 
             all_res = pi.execute()
 
         cookies = []
         set_cookies = []
 
-        for res in all_res:
-            if not res:
-                continue
+        with redis.utils.pipeline(self.redis) as pi:
+            for res, domain in zip(all_res, subds):
+                if not res:
+                    continue
+
+                for n, v in six.iteritems(res):
+                    n = n.decode('utf-8')
+                    v = v.decode('utf-8')
+                    full = n + '=' + v
+                    cookies.append(full.split(';')[0])
+                    set_cookies.append(('Set-Cookie', full + '; Max-Age=' + str(self.expire_time)))
+
+                pi.expire(cookie_key + '.' + domain, self.expire_time)
 
-            for n, v in six.iteritems(res):
-                n = n.decode('utf-8')
-                v = v.decode('utf-8')
-                full = n + '=' + v
-                cookies.append(full.split(';')[0])
-                set_cookies.append(('Set-Cookie', full + '; Max-Age=120'))
 
         cookies = ';'.join(cookies)
         return cookies, set_cookies
 
+    def add_cookie(self, cookie_key, domain, name, value):
+        if domain[0] != '.':
+            domain = '.' + domain
+
+        with redis.utils.pipeline(self.redis) as pi:
+            pi.hset(cookie_key + domain, name, value)
+            pi.expire(cookie_key + domain, self.expire_time)
+
     @staticmethod
     def get_subdomains(url):
         tld = tldextract.extract(url)
@@ -72,9 +85,9 @@ class CookieTracker(object):
 
 # =============================================================================
 class DomainCacheCookieRewriter(WbUrlBaseCookieRewriter):
-    def __init__(self, url_rewriter, redis, cookie_key):
+    def __init__(self, url_rewriter, cookie_tracker, cookie_key):
         super(DomainCacheCookieRewriter, self).__init__(url_rewriter)
-        self.redis = redis
+        self.cookie_tracker = cookie_tracker
         self.cookie_key = cookie_key
 
     def rewrite_cookie(self, name, morsel):
@@ -98,9 +111,10 @@ class DomainCacheCookieRewriter(WbUrlBaseCookieRewriter):
             if morsel.get('secure'):
                 string += '; Secure'
 
-            with redis.utils.pipeline(self.redis) as pi:
-                pi.hset(self.cookie_key + domain, morsel.key, string)
-                pi.expire(self.cookie_key + domain, 120)
+            self.cookie_tracker.add_cookie(self.cookie_key,
+                                           domain,
+                                           morsel.key,
+                                           string)
 
         # else set cookie to rewritten path
         if morsel.get('path'):