Handle custom 403/404 pages

staticshield.py

@@ -33,7 +33,7 @@ Session(app)
 # Verify the required configuration
 # SERVE_DIR: Base dir of the files we want to serve; Flask will take care not to escape this dir
 # MOTHERSHIP: Mothership server and login-url, which will redirect here with a sessionstart/SEKRIT
-config_vars = ['SERVE_DIR', 'MOTHERSHIP']
+config_vars = ['SERVE_DIR', 'MOTHERSHIP', 'ERROR_PAGES_DIR']
 for config_var in config_vars:
     if config_var not in app.config:
         print(f'FLASK_{config_var} env var not set')
@@ -42,6 +42,22 @@ for config_var in config_vars:
         app.logger.info('Config env %s with value "%s"', config_var, app.config[config_var])
 
 
+def handle_error_code(http_code):
+    """Handle 403/404/whatever HTTP response; either with a simple default, or with custom file.
+
+    :param int http_code: HTTP status code to handle and return
+    """
+    path = f'{http_code}.html'
+    if app.config['ERROR_PAGES_DIR']:
+        file_path = os.path.join(app.config['ERROR_PAGES_DIR'], path)
+        app.logger.debug(file_path)
+        if os.path.isfile(file_path):
+            app.logger.info('Serving file %s', str(file_path))
+            # This takes a base directory and a path, and ensures that the path is contained in the directory, which makes it safe to accept user-provided paths.
+            return send_from_directory(app.config['ERROR_PAGES_DIR'], path), http_code
+    return 'Unable to set up session', http_code
+
+
 @app.route('/<path:path>', methods=['GET', 'POST'])
 def all_routes(path):
     """Intercept all routes and proxy it through a session.
@@ -67,7 +83,7 @@ def all_routes(path):
 
         if secret == 'denied':
             # Mother says no
-            return 'Unable to set up session', 403
+            return handle_error_code(403)
 
         # Ask the mothership if the secret is known to them, to prevent someone from just making up a URL
         # Mothership will invalidate this secret token upon handling this request to prevent replay
@@ -87,16 +103,16 @@ def all_routes(path):
                     app.logger.warning('new session aborted, secret "%s" was incorrect, not redirecting to %s', secret, redirect_path)
                 except ValueError as e:
                     app.logger.error('Error while decoding challenge response: %s', str(e))
-                    return 'Unable to set up session', 403
+                    return handle_error_code(403)
         except urllib.error.URLError as e:
             app.logger.error('lolwtf, server not found: %s', str(e.reason))
-            return 'Unable to set up session', 403
+            return handle_error_code(403)
         except urllib.error.HTTPError as e:
             app.logger.error('lolnope, server says no: %s', str(e.reason))
-            return 'Unable to set up session', 403
+            return handle_error_code(403)
 
         # Fallback to disallow
-        return 'Unable to set up session', 403
+        return handle_error_code(403)
 
     # check if the users exist or not
     if not session.get('id'):
@@ -114,7 +130,7 @@ def all_routes(path):
         return send_from_directory(app.config['SERVE_DIR'], path)
     else:
         app.logger.error('File not found: %s', str(file_path))
-        return 'Sorry, 404'
+        return handle_error_code(404)
 
 
 @app.route('/')