From 407da7528ba70565da46dc6363745e962c1e9983 Mon Sep 17 00:00:00 2001
From: Ilya Kreymer <ikreymer@gmail.com>
Date: Thu, 31 Jul 2014 17:02:26 -0700
Subject: [PATCH] proxy/rewrite: don't rewrite headers banner_only

---
 pywb/framework/proxy.py         |  7 ++++++-
 pywb/rewrite/header_rewriter.py | 16 +++++++++++-----
 pywb/rewrite/rewrite_content.py |  4 ++++
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/pywb/framework/proxy.py b/pywb/framework/proxy.py
index 90ff2fd0..ab322374 100644
--- a/pywb/framework/proxy.py
+++ b/pywb/framework/proxy.py
@@ -181,7 +181,12 @@ class ProxyRouter(object):
         elif is_https:
             wbrequest.wb_url.mod = 'bn_'
 
-        return route.handler(wbrequest)
+        response = route.handler(wbrequest)
+
+        if wbrequest.wb_url and wbrequest.wb_url.is_replay():
+            response.status_headers.replace_header('Cache-Control', 'no-cache')
+
+        return response
 
     def get_request_socket(self, env):
         if not self.ca:
diff --git a/pywb/rewrite/header_rewriter.py b/pywb/rewrite/header_rewriter.py
index 2dfc824d..fd41eba8 100644
--- a/pywb/rewrite/header_rewriter.py
+++ b/pywb/rewrite/header_rewriter.py
@@ -37,7 +37,7 @@ class HeaderRewriter:
 
     ENCODING_HEADERS = ['content-encoding']
 
-    REMOVE_HEADERS = ['transfer-encoding']
+    REMOVE_HEADERS = ['transfer-encoding', 'content-security-policy']
 
     PROXY_NO_REWRITE_HEADERS = ['content-length']
 
@@ -90,7 +90,10 @@ class HeaderRewriter:
         new_headers = []
         removed_header_dict = {}
 
-        cookie_rewriter = urlrewriter.get_cookie_rewriter()
+        if urlrewriter:
+            cookie_rewriter = urlrewriter.get_cookie_rewriter()
+        else:
+            cookie_rewriter = None
 
         for (name, value) in headers:
 
@@ -99,7 +102,7 @@ class HeaderRewriter:
             if lowername in self.PROXY_HEADERS:
                 new_headers.append((name, value))
 
-            elif lowername in self.URL_REWRITE_HEADERS:
+            elif urlrewriter and lowername in self.URL_REWRITE_HEADERS:
                 new_headers.append((name, urlrewriter.rewrite(value)))
 
             elif lowername in self.ENCODING_HEADERS:
@@ -109,7 +112,8 @@ class HeaderRewriter:
                     new_headers.append((name, value))
 
             elif lowername in self.REMOVE_HEADERS:
-                    removed_header_dict[lowername] = value
+                removed_header_dict[lowername] = value
+                new_headers.append((self.header_prefix + name, value))
 
             elif (lowername in self.PROXY_NO_REWRITE_HEADERS and
                   not content_rewritten):
@@ -120,7 +124,9 @@ class HeaderRewriter:
                 cookie_list = cookie_rewriter.rewrite(value)
                 new_headers.extend(cookie_list)
 
-            else:
+            elif urlrewriter:
                 new_headers.append((self.header_prefix + name, value))
+            else:
+                new_headers.append((name, value))
 
         return (new_headers, removed_header_dict)
diff --git a/pywb/rewrite/rewrite_content.py b/pywb/rewrite/rewrite_content.py
index 93ec396b..e81fdf9a 100644
--- a/pywb/rewrite/rewrite_content.py
+++ b/pywb/rewrite/rewrite_content.py
@@ -63,6 +63,10 @@ class RewriteContent:
             status_headers, stream = self.sanitize_content(headers, stream)
             return (status_headers, self.stream_to_gen(stream), False)
 
+
+        if wb_url.is_banner_only:
+            urlrewriter = None
+
         (rewritten_headers, stream) = self.rewrite_headers(urlrewriter,
                                                            headers,
                                                            stream)