From 15ad56c024796cae6f61dd2a493262c638a3703f Mon Sep 17 00:00:00 2001
From: Ilya Kreymer <ikreymer@gmail.com>
Date: Tue, 21 Mar 2017 11:18:53 -0700
Subject: [PATCH] rewrite dash: support for using custom rewriting function
 (for FB) rewrite_fb_dash() added for rewriting dash xml, embedded in js,
 embedded in html todo: refactor to make more general support for custom
 rewriting functions regex_rewriter: add ':' to exclude from rewrite again

---
 pywb/rewrite/regex_rewriters.py | 13 ++++++++++++-
 pywb/rewrite/rewrite_dash.py    | 20 ++++++++++++++++++++
 pywb/rules.yaml                 |  4 ++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/pywb/rewrite/regex_rewriters.py b/pywb/rewrite/regex_rewriters.py
index b0adc383..d2395ee3 100644
--- a/pywb/rewrite/regex_rewriters.py
+++ b/pywb/rewrite/regex_rewriters.py
@@ -3,6 +3,15 @@ import re
 from pywb.rewrite.url_rewriter import UrlRewriter
 
 
+#=================================================================
+def load_function(string):
+    import importlib
+
+    string = string.split(':', 1)
+    mod = importlib.import_module(string[0])
+    return getattr(mod, string[1])
+
+
 #=================================================================
 class RegexRewriter(object):
     #@staticmethod
@@ -93,6 +102,8 @@ class RegexRewriter(object):
                 match = obj.get('match')
                 if 'rewrite' in obj:
                     replace = RegexRewriter.archival_rewrite(rewriter)
+                elif 'function' in obj:
+                    replace = load_function(obj['function'])
                 else:
                     replace = RegexRewriter.format(obj.get('replace', '{0}'))
                 group = obj.get('group', 0)
@@ -132,7 +143,7 @@ class JSLocationRewriterMixin(object):
     def __init__(self, rewriter, rules=[], prefix='WB_wombat_'):
         rules = rules + [
           #  (r'(?<![/$])\blocation\b(?!\":)', RegexRewriter.add_prefix(prefix), 0),
-          (r'(?<![$\'"])\b(?:location|top)\b(?![$\'"])', RegexRewriter.add_prefix(prefix), 0),
+          (r'(?<![$\'"])\b(?:location|top)\b(?![$\'":])', RegexRewriter.add_prefix(prefix), 0),
 
           (r'(?<=\.)postMessage\b\(', RegexRewriter.add_prefix('__WB_pmw(window).'), 0),
 
diff --git a/pywb/rewrite/rewrite_dash.py b/pywb/rewrite/rewrite_dash.py
index a8441dfa..c01a5eab 100644
--- a/pywb/rewrite/rewrite_dash.py
+++ b/pywb/rewrite/rewrite_dash.py
@@ -105,3 +105,23 @@ class RewriteDASHMixin(object):
         buff_io.seek(0)
         return buff_io
 
+
+def rewrite_fb_dash(string):
+    DASH_SPLIT = r'\n",dash_prefetched_representation_ids:'
+    inx = string.find(DASH_SPLIT)
+    if inx < 0:
+        return string
+
+    string = string[:inx]
+
+    buff = string.encode('utf-8').decode('unicode-escape')
+    buff = buff.encode('utf-8')
+    io = BytesIO(buff)
+    io, best_ids = RewriteDASHMixin.rewrite_dash(io)
+    string = json.dumps(io.read().decode('utf-8'))
+    string = string[1:-1].replace('<', r'\x3C')
+
+    string += DASH_SPLIT
+    string += json.dumps(best_ids)
+    return string
+
diff --git a/pywb/rules.yaml b/pywb/rules.yaml
index e731e10f..6ab37cf7 100644
--- a/pywb/rules.yaml
+++ b/pywb/rules.yaml
@@ -58,6 +58,10 @@ rules:
             - match: 'Bootloader\.configurePage.*?;'
               replace: '/* {0} */'
 
+            - match: 'dash_manifest:"(.*",dash_prefetched_representation_ids:.*?])'
+              group: 1
+              function: 'pywb.rewrite.rewrite_dash:rewrite_fb_dash'
+
         parse_comments: true
 
     - url_prefix: 'com,facebook'