Merge pull request #1358 from whew/add_security_headers

Add security headers to every response

5 files changed, 31 insertions, 54 deletions

diff --git a/cli/onionshare_cli/web/chat_mode.py b/cli/onionshare_cli/web/chat_mode.py
index f6dc2d1a..e92ce385 100644
--- a/cli/onionshare_cli/web/chat_mode.py
+++ b/cli/onionshare_cli/web/chat_mode.py
@@ -68,15 +68,12 @@ class ChatModeWeb:
             )
 
             self.web.add_request(self.web.REQUEST_LOAD, request.path)
-            r = make_response(
-                render_template(
+            return render_template(
                     "chat.html",
                     static_url_path=self.web.static_url_path,
                     username=session.get("name"),
                     title=self.web.settings.get("general", "title"),
-                )
             )
-            return self.web.add_security_headers(r)
 
         @self.web.app.route("/update-session-username", methods=["POST"], provide_automatic_options=False)
         def update_session_username():
@@ -112,7 +109,7 @@ class ChatModeWeb:
                         success=False,
                     )
                 )
-            return self.web.add_security_headers(r)
+            return r
 
         @self.web.socketio.on("joined", namespace="/chat")
         def joined(message):
diff --git a/cli/onionshare_cli/web/receive_mode.py b/cli/onionshare_cli/web/receive_mode.py
index 76abb0a8..6b106d37 100644
--- a/cli/onionshare_cli/web/receive_mode.py
+++ b/cli/onionshare_cli/web/receive_mode.py
@@ -86,16 +86,13 @@ class ReceiveModeWeb:
             )
 
             self.web.add_request(self.web.REQUEST_LOAD, request.path)
-            r = make_response(
-                render_template(
-                    "receive.html",
-                    static_url_path=self.web.static_url_path,
-                    disable_text=self.web.settings.get("receive", "disable_text"),
-                    disable_files=self.web.settings.get("receive", "disable_files"),
-                    title=self.web.settings.get("general", "title"),
-                )
+            return render_template(
+                "receive.html",
+                static_url_path=self.web.static_url_path,
+                disable_text=self.web.settings.get("receive", "disable_text"),
+                disable_files=self.web.settings.get("receive", "disable_files"),
+                title=self.web.settings.get("general", "title")
             )
-            return self.web.add_security_headers(r)
 
         @self.web.app.route("/upload", methods=["POST"], provide_automatic_options=False)
         def upload(ajax=False):
@@ -222,12 +219,11 @@ class ReceiveModeWeb:
                     )
                 else:
                     # It was the last upload and the timer ran out
-                    r = make_response(
+                    return make_response(
                         render_template("thankyou.html"),
                         static_url_path=self.web.static_url_path,
                         title=self.web.settings.get("general", "title"),
                     )
-                    return self.web.add_security_headers(r)
 
         @self.web.app.route("/upload-ajax", methods=["POST"], provide_automatic_options=False)
         def upload_ajax_public():
diff --git a/cli/onionshare_cli/web/send_base_mode.py b/cli/onionshare_cli/web/send_base_mode.py
index e448d2dd..27de598a 100644
--- a/cli/onionshare_cli/web/send_base_mode.py
+++ b/cli/onionshare_cli/web/send_base_mode.py
@@ -149,10 +149,9 @@ class SendBaseModeWeb:
 
         # If filesystem_path is None, this is the root directory listing
         files, dirs = self.build_directory_listing(path, filenames, filesystem_path)
-        r = self.directory_listing_template(
+        return self.directory_listing_template(
             path, files, dirs, breadcrumbs, breadcrumbs_leaf
         )
-        return self.web.add_security_headers(r)
 
     def build_directory_listing(self, path, filenames, filesystem_path):
         files = []
@@ -286,7 +285,6 @@ class SendBaseModeWeb:
             "filename*": "UTF-8''%s" % url_quote(basename),
         }
         r.headers.set("Content-Disposition", "inline", **filename_dict)
-        r = self.web.add_security_headers(r)
         (content_type, _) = mimetypes.guess_type(basename, strict=False)
         if content_type is not None:
             r.headers.set("Content-Type", content_type)
diff --git a/cli/onionshare_cli/web/share_mode.py b/cli/onionshare_cli/web/share_mode.py
index 51ddd674..1dd9af13 100644
--- a/cli/onionshare_cli/web/share_mode.py
+++ b/cli/onionshare_cli/web/share_mode.py
@@ -149,8 +149,7 @@ class ShareModeWeb(SendBaseModeWeb):
                 and self.download_in_progress
             )
             if deny_download:
-                r = make_response(render_template("denied.html"))
-                return self.web.add_security_headers(r)
+                return render_template("denied.html")
 
             # If download is allowed to continue, serve download page
             if self.should_use_gzip():
@@ -172,8 +171,7 @@ class ShareModeWeb(SendBaseModeWeb):
                 and self.download_in_progress
             )
             if deny_download:
-                r = make_response(render_template("denied.html"))
-                return self.web.add_security_headers(r)
+                return render_template("denied.html")
 
             # Prepare some variables to use inside generate() function below
             # which is outside of the request context
@@ -232,7 +230,6 @@ class ShareModeWeb(SendBaseModeWeb):
                 "filename*": "UTF-8''%s" % url_quote(basename),
             }
             r.headers.set("Content-Disposition", "attachment", **filename_dict)
-            r = self.web.add_security_headers(r)
             # guess content type
             (content_type, _) = mimetypes.guess_type(basename, strict=False)
             if content_type is not None:
diff --git a/cli/onionshare_cli/web/web.py b/cli/onionshare_cli/web/web.py
index 0f2dfe7e..3ba27ef7 100644
--- a/cli/onionshare_cli/web/web.py
+++ b/cli/onionshare_cli/web/web.py
@@ -191,6 +191,21 @@ class Web:
         Common web app routes between all modes.
         """
 
+        @self.app.after_request
+        def add_security_headers(r):
+            """
+            Add security headers to a response
+            """
+            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":
+                r.headers.set(
+                    "Content-Security-Policy",
+                    "default-src 'self'; frame-ancestors 'none'; form-action 'self'; base-uri 'self'; img-src 'self' data:;",
+                )
+            return r
+
         @self.app.errorhandler(404)
         def not_found(e):
             mode = self.get_mode()
@@ -232,10 +247,7 @@ class Web:
 
     def error403(self):
         self.add_request(Web.REQUEST_OTHER, request.path)
-        r = make_response(
-            render_template("403.html", static_url_path=self.static_url_path), 403
-        )
-        return self.add_security_headers(r)
+        return render_template("403.html", static_url_path=self.static_url_path), 403
 
     def error404(self, history_id):
         mode = self.get_mode()
@@ -247,10 +259,7 @@ class Web:
             )
 
         self.add_request(Web.REQUEST_OTHER, request.path)
-        r = make_response(
-            render_template("404.html", static_url_path=self.static_url_path), 404
-        )
-        return self.add_security_headers(r)
+        return render_template("404.html", static_url_path=self.static_url_path), 404
 
     def error405(self, history_id):
         mode = self.get_mode()
@@ -262,10 +271,7 @@ class Web:
             )
 
         self.add_request(Web.REQUEST_OTHER, request.path)
-        r = make_response(
-            render_template("405.html", static_url_path=self.static_url_path), 405
-        )
-        return self.add_security_headers(r)
+        return render_template("405.html", static_url_path=self.static_url_path), 405
 
     def error500(self, history_id):
         mode = self.get_mode()
@@ -277,24 +283,7 @@ class Web:
             )
 
         self.add_request(Web.REQUEST_OTHER, request.path)
-        r = make_response(
-            render_template("500.html", static_url_path=self.static_url_path), 500
-        )
-        return self.add_security_headers(r)
-
-    def add_security_headers(self, r):
-        """
-        Add security headers to a request
-        """
-        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":
-            r.headers.set(
-                "Content-Security-Policy",
-                "default-src 'self'; frame-ancestors 'none'; form-action 'self'; base-uri 'self'; img-src 'self' data:;",
-            )
-        return r
+        return render_template("500.html", static_url_path=self.static_url_path), 500
 
     def _safe_select_jinja_autoescape(self, filename):
         if filename is None: