From 3101e567f32d62ea05a86affdbc401004dd6a1be Mon Sep 17 00:00:00 2001
From: Ilya Kreymer <ikreymer@gmail.com>
Date: Tue, 3 Apr 2018 19:05:01 -0700
Subject: [PATCH] config: add support for forcing a scheme for url rewriting,
 eg: 'force_scheme: https', fixes #314

---
 pywb/apps/rewriterapp.py  |  6 +++++
 tests/test_force_https.py | 48 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+)
 create mode 100644 tests/test_force_https.py

diff --git a/pywb/apps/rewriterapp.py b/pywb/apps/rewriterapp.py
index b23e948e..cfa6e3e9 100644
--- a/pywb/apps/rewriterapp.py
+++ b/pywb/apps/rewriterapp.py
@@ -103,6 +103,8 @@ class RewriterApp(object):
         else:
             self.csp_header = None
 
+        self.force_scheme = config.get('force_scheme')
+
     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)
@@ -202,6 +204,10 @@ class RewriterApp(object):
     def render_content(self, wb_url, kwargs, environ):
         wb_url = wb_url.replace('#', '%23')
         wb_url = WbUrl(wb_url)
+
+        if self.force_scheme:
+            environ['wsgi.url_scheme'] = self.force_scheme
+
         is_timegate = self._check_accept_dt(wb_url, environ)
 
         host_prefix = self.get_host_prefix(environ)
diff --git a/tests/test_force_https.py b/tests/test_force_https.py
new file mode 100644
index 00000000..bac42758
--- /dev/null
+++ b/tests/test_force_https.py
@@ -0,0 +1,48 @@
+from .base_config_test import BaseConfigTest, fmod
+
+
+# ============================================================================
+class TestForceHttps(BaseConfigTest):
+    @classmethod
+    def setup_class(cls):
+        super(TestForceHttps, cls).setup_class('config_test.yaml',
+                                               custom_config={'force_scheme': 'https'})
+
+    def test_force_https_replay_1(self, fmod):
+        resp = self.get('/pywb/20140128051539{0}/http://example.com/', fmod)
+
+        assert '"https://localhost:80/pywb/20140128051539{0}/http://www.iana.org/domains/example"'.format(fmod) in resp.text, resp.text
+
+
+# ============================================================================
+class TestForceHttpsRedirect(BaseConfigTest):
+    @classmethod
+    def setup_class(cls):
+        super(TestForceHttpsRedirect, cls).setup_class('config_test_redirect_classic.yaml',
+                                                       custom_config={'force_scheme': 'https'})
+
+    def test_force_https_redirect_replay_1(self, fmod):
+        resp = self.get('/pywb/20140128051539{0}/http://example.com/', fmod)
+
+        assert resp.headers['Location'] == 'https://localhost:80/pywb/20140127171251{0}/http://example.com'.format(fmod)
+        resp = resp.follow()
+
+        assert resp.headers['Location'] == 'https://localhost:80/pywb/20140127171251{0}/http://example.com/'.format(fmod)
+        resp = resp.follow()
+
+        assert '"https://localhost:80/pywb/20140127171251{0}/http://www.iana.org/domains/example"'.format(fmod) in resp.text, resp.text
+
+
+# ============================================================================
+class TestForceHttpsRoot(BaseConfigTest):
+    @classmethod
+    def setup_class(cls):
+        super(TestForceHttpsRoot, cls).setup_class('config_test_root_coll.yaml',
+                                                   custom_config={'force_scheme': 'https'})
+
+    def test_force_https_root_replay_1(self, fmod):
+        resp = self.get('/20140128051539{0}/http://www.iana.org/domains/example', fmod)
+
+        assert resp.headers['Location'] == 'https://localhost:80/20140128051539{0}/https://www.iana.org/domains/reserved'.format(fmod)
+
+