Support sending a custom Content-Security-Policy header in Website mode

6 files changed, 103 insertions, 10 deletions

diff --git a/cli/onionshare_cli/__init__.py b/cli/onionshare_cli/__init__.py
index 060c5628..ded67ed6 100644
--- a/cli/onionshare_cli/__init__.py
+++ b/cli/onionshare_cli/__init__.py
@@ -150,7 +150,13 @@ def main(cwd=None):
         action="store_true",
         dest="disable_csp",
         default=False,
-        help="Publish website: Disable Content Security Policy header (allows your website to use third-party resources)",
+        help="Publish website: Disable the default Content Security Policy header (allows your website to use third-party resources)",
+    )
+    parser.add_argument(
+        "--custom_csp",
+        metavar="custom_csp",
+        default=None,
+        help="Publish website: Set a custom Content Security Policy header",
     )
     # Other
     parser.add_argument(
@@ -189,6 +195,7 @@ def main(cwd=None):
     disable_text = args.disable_text
     disable_files = args.disable_files
     disable_csp = bool(args.disable_csp)
+    custom_csp = args.custom_csp
     verbose = bool(args.verbose)
 
     # Verbose mode?
@@ -234,7 +241,15 @@ def main(cwd=None):
             mode_settings.set("receive", "disable_text", disable_text)
             mode_settings.set("receive", "disable_files", disable_files)
         if mode == "website":
-            mode_settings.set("website", "disable_csp", disable_csp)
+            if disable_csp and custom_csp:
+                print("You cannot disable the CSP and set a custom one. Either set --disable-csp or --custom-csp but not both.")
+                sys.exit()
+            if disable_csp:
+                mode_settings.set("website", "disable_csp", True)
+                mode_settings.set("website", "custom_csp", None)
+            if custom_csp:
+                mode_settings.set("website", "custom_csp", custom_csp)
+                mode_settings.set("website", "disable_csp", False)
     else:
         # See what the persistent mode was
         mode = mode_settings.get("persistent", "mode")
diff --git a/cli/onionshare_cli/mode_settings.py b/cli/onionshare_cli/mode_settings.py
index 47ff1c63..b94b1d25 100644
--- a/cli/onionshare_cli/mode_settings.py
+++ b/cli/onionshare_cli/mode_settings.py
@@ -55,7 +55,11 @@ class ModeSettings:
                 "disable_text": False,
                 "disable_files": False,
             },
-            "website": {"disable_csp": False, "filenames": []},
+            "website": {
+                "disable_csp": False,
+                "custom_csp": None,
+                "filenames": []
+            },
             "chat": {"room": "default"},
         }
         self._settings = {}
diff --git a/cli/onionshare_cli/web/web.py b/cli/onionshare_cli/web/web.py
index e12fccc7..e0cf97f3 100644
--- a/cli/onionshare_cli/web/web.py
+++ b/cli/onionshare_cli/web/web.py
@@ -199,11 +199,18 @@ class Web:
             for header, value in self.security_headers:
                 r.headers.set(header, value)
             # Set a CSP header unless in website mode and the user has disabled it
-            if not self.settings.get("website", "disable_csp") or self.mode != "website":
+            default_csp = "default-src 'self'; frame-ancestors 'none'; form-action 'self'; base-uri 'self'; img-src 'self' data:;"
+            if self.mode != "website" or (not self.settings.get("website", "disable_csp") and not self.settings.get("website", "custom_csp")):
                 r.headers.set(
                     "Content-Security-Policy",
-                    "default-src 'self'; frame-ancestors 'none'; form-action 'self'; base-uri 'self'; img-src 'self' data:;",
+                    default_csp
                 )
+            else:
+                if self.settings.get("website", "custom_csp"):
+                    r.headers.set(
+                        "Content-Security-Policy",
+                        self.settings.get("website", "custom_csp")
+                    )
             return r
 
         @self.app.errorhandler(404)
diff --git a/desktop/src/onionshare/resources/locale/en.json b/desktop/src/onionshare/resources/locale/en.json
index d405c702..f8c4cd2b 100644
--- a/desktop/src/onionshare/resources/locale/en.json
+++ b/desktop/src/onionshare/resources/locale/en.json
@@ -203,7 +203,8 @@
     "mode_settings_receive_disable_text_checkbox": "Disable submitting text",
     "mode_settings_receive_disable_files_checkbox": "Disable uploading files",
     "mode_settings_receive_webhook_url_checkbox": "Use notification webhook",
-    "mode_settings_website_disable_csp_checkbox": "Don't send Content Security Policy header (allows your website to use third-party resources)",
+    "mode_settings_website_disable_csp_checkbox": "Don't send default Content Security Policy header (allows your website to use third-party resources)",
+    "mode_settings_website_custom_csp_checkbox": "Send a custom Content Security Policy header",
     "gui_all_modes_transfer_finished_range": "Transferred {} - {}",
     "gui_all_modes_transfer_finished": "Transferred {}",
     "gui_all_modes_transfer_canceled_range": "Canceled {} - {}",
@@ -232,4 +233,4 @@
     "moat_captcha_error": "The solution is not correct. Please try again.",
     "moat_solution_empty_error": "You must enter the characters from the image",
     "mode_tor_not_connected_label": "OnionShare is not connected to the Tor network"
-}
\ No newline at end of file
+}
diff --git a/desktop/src/onionshare/tab/mode/website_mode/__init__.py b/desktop/src/onionshare/tab/mode/website_mode/__init__.py
index 73c4bad2..0acbc1a2 100644
--- a/desktop/src/onionshare/tab/mode/website_mode/__init__.py
+++ b/desktop/src/onionshare/tab/mode/website_mode/__init__.py
@@ -49,6 +49,7 @@ class WebsiteMode(Mode):
         self.web = Web(self.common, True, self.settings, "website")
 
         # Settings
+        # Disable CSP option
         self.disable_csp_checkbox = QtWidgets.QCheckBox()
         self.disable_csp_checkbox.clicked.connect(self.disable_csp_checkbox_clicked)
         self.disable_csp_checkbox.setText(
@@ -63,6 +64,26 @@ class WebsiteMode(Mode):
             self.disable_csp_checkbox
         )
 
+        # Custom CSP option
+        self.custom_csp_checkbox = QtWidgets.QCheckBox()
+        self.custom_csp_checkbox.clicked.connect(self.custom_csp_checkbox_clicked)
+        self.custom_csp_checkbox.setText(strings._("mode_settings_website_custom_csp_checkbox"))
+        if self.settings.get("website", "custom_csp") and not self.settings.get("website", "disable_csp"):
+            self.custom_csp_checkbox.setCheckState(QtCore.Qt.Checked)
+        else:
+            self.custom_csp_checkbox.setCheckState(QtCore.Qt.Unchecked)
+        self.custom_csp = QtWidgets.QLineEdit()
+        self.custom_csp.setPlaceholderText(
+            "default-src 'self'; frame-ancestors 'none'; form-action 'self'; base-uri 'self'; img-src 'self' data:;"
+        )
+        self.custom_csp.editingFinished.connect(self.custom_csp_editing_finished)
+
+        custom_csp_layout = QtWidgets.QHBoxLayout()
+        custom_csp_layout.setContentsMargins(0, 0, 0, 0)
+        custom_csp_layout.addWidget(self.custom_csp_checkbox)
+        custom_csp_layout.addWidget(self.custom_csp)
+        self.mode_settings_widget.mode_specific_layout.addLayout(custom_csp_layout)
+
         # File selection
         self.file_selection = FileSelection(
             self.common,
@@ -181,11 +202,42 @@ class WebsiteMode(Mode):
 
     def disable_csp_checkbox_clicked(self):
         """
-        Save disable CSP setting to the tab settings
+        Save disable CSP setting to the tab settings. Uncheck 'custom CSP'
+        setting if disabling CSP altogether.
         """
         self.settings.set(
             "website", "disable_csp", self.disable_csp_checkbox.isChecked()
         )
+        if self.disable_csp_checkbox.isChecked():
+            self.custom_csp_checkbox.setCheckState(QtCore.Qt.Unchecked)
+            self.custom_csp_checkbox.setEnabled(False)
+        else:
+            self.custom_csp_checkbox.setEnabled(True)
+
+    def custom_csp_checkbox_clicked(self):
+        """
+        Uncheck 'disable CSP' setting if custom CSP is used.
+        """
+        if self.custom_csp_checkbox.isChecked():
+            self.disable_csp_checkbox.setCheckState(QtCore.Qt.Unchecked)
+            self.disable_csp_checkbox.setEnabled(False)
+            self.settings.set(
+                "website", "custom_csp", self.custom_csp
+            )
+        else:
+            self.disable_csp_checkbox.setEnabled(True)
+            self.custom_csp.setText("")
+            self.settings.set(
+                "website", "custom_csp", None
+            )
+
+    def custom_csp_editing_finished(self):
+        if self.custom_csp.text().strip() == "":
+            self.custom_csp.setText("")
+            self.settings.set("website", "custom_csp", None)
+        else:
+            custom_csp = self.custom_csp.text()
+            self.settings.set("website", "custom_csp", custom_csp)
 
     def get_stop_server_autostop_timer_text(self):
         """
diff --git a/desktop/tests/test_gui_website.py b/desktop/tests/test_gui_website.py
index e736874a..a46e21a9 100644
--- a/desktop/tests/test_gui_website.py
+++ b/desktop/tests/test_gui_website.py
@@ -22,8 +22,10 @@ class TestWebsite(GuiBaseTest):
         QtTest.QTest.qWait(500, self.gui.qtapp)
         if tab.settings.get("website", "disable_csp"):
             self.assertFalse("Content-Security-Policy" in r.headers)
+        elif tab.settings.get("website", "custom_csp"):
+            self.assertEqual(tab.settings.get("website", "custom_csp"), r.headers["Content-Security-Policy"])
         else:
-            self.assertTrue("Content-Security-Policy" in r.headers)
+            self.assertEqual("default-src 'self'; frame-ancestors 'none'; form-action 'self'; base-uri 'self'; img-src 'self' data:;", r.headers["Content-Security-Policy"])
 
     def run_all_website_mode_setup_tests(self, tab):
         """Tests in website mode prior to starting a share"""
@@ -77,12 +79,24 @@ class TestWebsite(GuiBaseTest):
         self.run_all_website_mode_download_tests(tab)
         self.close_all_tabs()
 
-    def test_csp_enabled(self):
+    def test_csp_disabled(self):
         """
         Test disabling CSP
         """
         tab = self.new_website_tab()
         tab.get_mode().disable_csp_checkbox.click()
+        self.assertFalse(tab.get_mode().custom_csp_checkbox.isEnabled())
+        self.run_all_website_mode_download_tests(tab)
+        self.close_all_tabs()
+
+    def test_csp_custom(self):
+        """
+        Test a custom CSP
+        """
+        tab = self.new_website_tab()
+        tab.get_mode().custom_csp_checkbox.click()
+        self.assertFalse(tab.get_mode().disable_csp_checkbox.isEnabled())
+        tab.settings.set("website", "custom_csp", "default-src 'self'")
         self.run_all_website_mode_download_tests(tab)
         self.close_all_tabs()