From 72fe274e066f89100053735aa3c75a58cb1ca770 Mon Sep 17 00:00:00 2001
From: Ilya Kreymer <ikreymer@gmail.com>
Date: Mon, 4 Aug 2014 18:41:33 -0700
Subject: [PATCH] rewrite: better parametrize live rewrite timestamp cookie,
 restore LiveResourceException

---
 pywb/static/wb_frame.js             | 13 ++++++++-----
 pywb/webapp/live_rewrite_handler.py | 18 +++++++++++++-----
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/pywb/static/wb_frame.js b/pywb/static/wb_frame.js
index 0d7be28d..2dd88228 100644
--- a/pywb/static/wb_frame.js
+++ b/pywb/static/wb_frame.js
@@ -1,5 +1,8 @@
 var update_wb_url = push_state;
 
+var LIVE_COOKIE_REGEX = /pywb.timestamp=([\d]{1,14})/;
+
+
 function make_outer_url(url, ts)
 {
     if (ts) {
@@ -88,8 +91,7 @@ window.onpopstate = function(event) {
 }
 
 function extract_ts_cookie(value) {
-    var regex = /pywb.timestamp=([\d]{1,14})/;
-    var result = value.match(regex);
+    var result = value.match(LIVE_COOKIE_REGEX);
     if (result) {
         return result[1];
     } else {
@@ -114,10 +116,11 @@ function iframe_loaded(event) {
         ts = iframe.wbinfo.timestamp;
         is_live = iframe.wbinfo.is_live;
     } else {
-        ts = extract_ts(iframe.location.href);
-        if (!ts) {
+        ts = extract_ts_cookie(iframe.document.cookie);
+        if (ts) {
             is_live = true;
-            ts = extract_ts_cookie(iframe.document.cookie);
+        } else {
+            ts = extract_ts(iframe.location.href);
         }
     }
     capture_str = _wb_js.ts_to_date(ts, true);
diff --git a/pywb/webapp/live_rewrite_handler.py b/pywb/webapp/live_rewrite_handler.py
index 15b9b3f8..a1b602d4 100644
--- a/pywb/webapp/live_rewrite_handler.py
+++ b/pywb/webapp/live_rewrite_handler.py
@@ -19,6 +19,9 @@ class LiveResourceException(WbException):
 
 #=================================================================
 class RewriteHandler(SearchPageWbUrlHandler):
+
+    LIVE_COOKIE = 'pywb.timestamp={0}; max-age=60';
+
     def __init__(self, config):
         super(RewriteHandler, self).__init__(config)
 
@@ -28,6 +31,8 @@ class RewriteHandler(SearchPageWbUrlHandler):
 
         self.head_insert_view = HeadInsertView.init_from_config(config)
 
+        self.live_cookie = config.get('live-cookie', self.LIVE_COOKIE)
+
     def handle_request(self, wbrequest):
         try:
             return self.render_content(wbrequest)
@@ -35,8 +40,7 @@ class RewriteHandler(SearchPageWbUrlHandler):
         except Exception as exc:
             url = wbrequest.wb_url.url
             msg = 'Could not load the url from the live web: ' + url
-            #raise LiveResourceException(msg=msg, url=url)
-            raise
+            raise LiveResourceException(msg=msg, url=url)
 
     def _live_request_headers(self, wbrequest):
         return {}
@@ -58,9 +62,13 @@ class RewriteHandler(SearchPageWbUrlHandler):
         return self._make_response(wbrequest, *result)
 
     def _make_response(self, wbrequest, status_headers, gen, is_rewritten):
-        cdx = wbrequest.env['pywb.cdx']
-        cookie = 'pywb.timestamp=' + cdx['timestamp'] + '; max-age=60'
-        status_headers.headers.append(('Set-Cookie', cookie))
+        # if cookie set, pass recorded timestamp info via cookie
+        # so that client side may be able to access it
+        # used by framed mode to update frame banner
+        if self.live_cookie:
+            cdx = wbrequest.env['pywb.cdx']
+            value = self.live_cookie.format(cdx['timestamp'])
+            status_headers.headers.append(('Set-Cookie', value))
 
         return WbResponse(status_headers, gen)