From 902f6659f490cc1d29f68f10ef163a708719c72c Mon Sep 17 00:00:00 2001
From: Ilya Kreymer <ikreymer@gmail.com>
Date: Thu, 5 Oct 2017 19:59:37 -0700
Subject: [PATCH] rewriterapp: add default csp header, overridable via
 'csp-header' config setting

---
 pywb/apps/rewriterapp.py  | 16 ++++++++++++++++
 tests/test_integration.py |  5 +++++
 2 files changed, 21 insertions(+)

diff --git a/pywb/apps/rewriterapp.py b/pywb/apps/rewriterapp.py
index 1b846bcf..f90b5a19 100644
--- a/pywb/apps/rewriterapp.py
+++ b/pywb/apps/rewriterapp.py
@@ -48,6 +48,8 @@ class UpstreamException(WbException):
 class RewriterApp(object):
     VIDEO_INFO_CONTENT_TYPE = 'application/vnd.youtube-dl_formats+json'
 
+    DEFAULT_CSP = "default-src 'unsafe-eval' 'unsafe-inline' 'self' data: blob: mediastream: ws: wss: ; form-action 'self'"
+
     def __init__(self, framed_replay=False, jinja_env=None, config=None, paths=None):
         self.loader = ArcWarcRecordLoader()
 
@@ -89,6 +91,16 @@ class RewriterApp(object):
 
         self.enable_memento = self.config.get('enable_memento')
 
+        csp_header = self.config.get('csp-header', self.DEFAULT_CSP)
+        if csp_header:
+            self.csp_header = ('Content-Security-Policy', csp_header)
+        else:
+            self.csp_header = None
+
+    def add_csp_header(self, wb_url, status_headers):
+        if self.csp_header and wb_url.mod == self.replay_mod:
+            status_headers.headers.append(self.csp_header)
+
     def _html_templ(self, name):
         value = self.config.get(name)
         if not value:
@@ -302,6 +314,10 @@ class RewriterApp(object):
         if set_content_loc:
             status_headers.headers.append(('Content-Location', urlrewriter.get_new_url(timestamp=cdx['timestamp'],
                                                                                        url=cdx['url'])))
+
+        if not is_proxy:
+            self.add_csp_header(wb_url, status_headers)
+
         response = WbResponse(status_headers, gen)
 
         return response
diff --git a/tests/test_integration.py b/tests/test_integration.py
index 5d30519b..6e0fe345 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -78,6 +78,8 @@ class TestWbIntegration(BaseConfigTest):
         assert '"20140127171238"' in resp.text
         assert 'http://www.iana.org/' in resp.text, resp.text
 
+        assert 'Content-Security-Policy' not in resp.headers
+
     def test_replay_content(self, fmod):
         resp = self.get('/pywb/20140127171238{0}/http://www.iana.org/', fmod)
         self._assert_basic_html(resp)
@@ -89,6 +91,9 @@ class TestWbIntegration(BaseConfigTest):
 
         assert ('wbinfo.is_framed = ' + ('true' if fmod else 'false')) in resp.text
 
+        csp = "default-src 'unsafe-eval' 'unsafe-inline' 'self' data: blob: mediastream: ws: wss: ; form-action 'self'"
+        assert resp.headers['Content-Security-Policy'] == csp
+
     def test_replay_fuzzy_1(self, fmod):
         resp = self.get('/pywb/20140127171238{0}/http://www.iana.org/?_=123', fmod)
         assert resp.status_int == 200