From 40b0a291a9eaea09b02fc7009850c03d15d7e629 Mon Sep 17 00:00:00 2001
From: Ilya Kreymer <ikreymer@gmail.com>
Date: Thu, 20 Oct 2016 11:30:14 -0700
Subject: [PATCH] rewrite: don't rewrite ajax-requested html content js regex:
 add special regex to rewrite '?location:'

---
 pywb/rewrite/regex_rewriters.py     | 2 ++
 pywb/rewrite/rewrite_amf.py         | 4 ++--
 pywb/rewrite/rewrite_content.py     | 8 ++++++--
 pywb/webapp/live_rewrite_handler.py | 3 +++
 pywb/webapp/replay_views.py         | 3 +++
 5 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/pywb/rewrite/regex_rewriters.py b/pywb/rewrite/regex_rewriters.py
index 4d2747b3..f17b5f24 100644
--- a/pywb/rewrite/regex_rewriters.py
+++ b/pywb/rewrite/regex_rewriters.py
@@ -134,6 +134,8 @@ class JSLocationRewriterMixin(object):
           #  (r'(?<![/$])\blocation\b(?!\":)', RegexRewriter.add_prefix(prefix), 0),
           (r'(?<![/$\'"-])\b(?:location|top)\b(?!(?:\":|:|=\d|-))', RegexRewriter.add_prefix(prefix), 0),
 
+          (r'(?<=[?])location[:]', RegexRewriter.add_prefix(prefix), 0),
+
           (r'(?<=\.)postMessage\b\(', RegexRewriter.add_prefix('__WB_pmw(window).'), 0),
 
           (r'(?<=\.)frameElement\b', RegexRewriter.add_prefix(prefix), 0),
diff --git a/pywb/rewrite/rewrite_amf.py b/pywb/rewrite/rewrite_amf.py
index 0cfd217f..9fcae784 100644
--- a/pywb/rewrite/rewrite_amf.py
+++ b/pywb/rewrite/rewrite_amf.py
@@ -6,12 +6,12 @@ from pywb.rewrite.rewrite_content import RewriteContent
 # ============================================================================
 # Expiermental: not fully tested
 class RewriteContentAMF(RewriteContent):  #pragma: no cover
-    def handle_custom_rewrite(self, rewritten_headers, stream, mod, env):
+    def handle_custom_rewrite(self, rewritten_headers, stream, urlrewriter, mod, env):
         if rewritten_headers.status_headers.get_header('Content-Type') == 'application/x-amf':
             stream = self.rewrite_amf(stream, env)
 
         return (super(RewriteContentAMF, self).
-                handle_custom_rewrite(rewritten_headers, stream, mod, env))
+                handle_custom_rewrite(rewritten_headers, stream, urlrewriter, mod, env))
 
     def rewrite_amf(self, stream, env):
         try:
diff --git a/pywb/rewrite/rewrite_content.py b/pywb/rewrite/rewrite_content.py
index 3eb47451..a1186f38 100644
--- a/pywb/rewrite/rewrite_content.py
+++ b/pywb/rewrite/rewrite_content.py
@@ -120,6 +120,7 @@ class RewriteContent(object):
 
         res = self.handle_custom_rewrite(rewritten_headers,
                                          stream,
+                                         urlrewriter,
                                          wb_url.mod,
                                          env)
         if res:
@@ -255,7 +256,9 @@ class RewriteContent(object):
 
         return (status_headers, gen, True)
 
-    def handle_custom_rewrite(self, rewritten_headers, stream, mod, env):
+    def handle_custom_rewrite(self, rewritten_headers, stream,
+                              urlrewriter, mod, env):
+
         text_type = rewritten_headers.text_type
         status_headers = rewritten_headers.status_headers
 
@@ -263,7 +266,8 @@ class RewriteContent(object):
         if text_type is None:
             return (status_headers, self.stream_to_gen(stream), False)
 
-        if text_type == 'plain' and not mod in ('js_', 'cs_'):
+        if ((text_type == 'html' and urlrewriter.rewrite_opts.get('ajax')) or
+            (text_type == 'plain' and not mod in ('js_', 'cs_'))):
             rewritten_headers.readd_rewrite_removed()
             return (status_headers, self.stream_to_gen(stream), False)
 
diff --git a/pywb/webapp/live_rewrite_handler.py b/pywb/webapp/live_rewrite_handler.py
index 38014076..56143583 100644
--- a/pywb/webapp/live_rewrite_handler.py
+++ b/pywb/webapp/live_rewrite_handler.py
@@ -48,6 +48,9 @@ class RewriteHandler(SearchPageWbUrlHandler):
             url = wbrequest.urlrewriter.get_new_url(type=type_, timestamp='')
             return WbResponse.redir_response(url)
 
+        if wbrequest.options['is_ajax']:
+            wbrequest.urlrewriter.rewrite_opts['is_ajax'] = True
+
         try:
             return self.render_content(wbrequest)
 
diff --git a/pywb/webapp/replay_views.py b/pywb/webapp/replay_views.py
index 3637b152..69e9ea2d 100644
--- a/pywb/webapp/replay_views.py
+++ b/pywb/webapp/replay_views.py
@@ -163,6 +163,9 @@ class ReplayView(object):
         if wbrequest and wbrequest.wb_url:
             wbrequest.wb_url.url = cdx['url']
 
+        if wbrequest.options['is_ajax']:
+            wbrequest.urlrewriter.rewrite_opts['is_ajax'] = True
+
         head_insert_func = None
         if self.head_insert_view:
             head_insert_func = (self.head_insert_view.