Document the error pages dir

README.md

@@ -10,7 +10,11 @@ Create a file (e.g., called `run.sh`) with the following:
 #!/bin/bash
 export FLASK_SERVE_DIR="/home/YOURUSER/workspace/somesite/build/html"
 export FLASK_MOTHERSHIP="http://localhost:8888/api/staticshield"
+# Optional path to 403.html, 404.html to show on those errors; leave empty to use default messages
+export FLASK_ERROR_PAGES_DIR=""
+#export FLASK_ERROR_PAGES_DIR="/home/YOURUSER/workspace/errorpages/"
 export FLASK_SESSION_COOKIE_NAME="staticshield"
+export FLASK_PERMANENT_SESSION_LIFETIME=7200
 
 flask --app staticshield run
 ```
@@ -25,7 +29,7 @@ ruff check --fix --select I .
 
 ## Deploying
 
-Create a virtualenv with flask, gunicorn:
+Create a virtualenv with Flask and gunicorn to run it:
 
 ```bash
 # Example, create wherever you like
@@ -54,6 +58,9 @@ WorkingDirectory=/srv/staticshield
 #StandardOutput=file:/srv/logs/staticshield.log
 Environment=FLASK_SERVE_DIR="/srv/some_static_website/html"
 Environment=FLASK_MOTHERSHIP="https://api.example.com/api/staticshield"
+# Optional path to 403.html, 404.html to show on those errors; leave empty to use default messages
+Environment=FLASK_ERROR_PAGES_DIR=""
+#Environment=FLASK_ERROR_PAGES_DIR="/srv/shared/errorpages/"
 Environment=FLASK_SESSION_COOKIE_NAME="staticshield"
 # Max session length of 2h
 Environment=FLASK_PERMANENT_SESSION_LIFETIME=7200
@@ -132,6 +139,11 @@ server {
         proxy_pass http://127.0.0.1:8000;
     }
 
+    # Optional, when using the FLASK_ERROR_PAGES_DIR
+    # location /404images/ {
+    #     alias /srv/shared/errorpages/;
+    # }
+
     location /favicon.ico {
         alias /srv/whatever/_static/favicon.ico;
     }

pyproject.toml

@@ -50,7 +50,7 @@ select = [
 "__init__.py" = ["F401"]
 
 [tool.ruff.lint.mccabe]
-max-complexity = 10
+max-complexity = 12
 
 [tool.ruff.lint.pycodestyle]
 max-doc-length = 180

staticshield.py

@@ -33,7 +33,8 @@ 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']
+# ERROR_PAGES_DIR: (optional) Base dir of 403.html, 404.html to show as alternative to default short error message
+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,8 +43,25 @@ 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('/')
 @app.route('/<path:path>', methods=['GET', 'POST'])
-def all_routes(path):
+def all_routes(path = ''):
     """Intercept all routes and proxy it through a session.
 
     Loosely based on https://stackoverflow.com/a/20648053
@@ -56,6 +74,8 @@ def all_routes(path):
 
         # The path we should have gotten back is of the format:
         # /sessionstart/SEKRIT_TOKEN/<the_url_to_redirect_on_here_afterwards>
+        # or, if the request is denied by the mothership:
+        # /sessionstart/denied
         secret_and_redirect = path.split('sessionstart/')[1]
         secret_redirect_split = secret_and_redirect.split('/')
         secret = secret_redirect_split[0]
@@ -63,6 +83,10 @@ def all_routes(path):
         if len(secret_redirect_split) > 1:
             redirect_path = '/'.join(secret_redirect_split[1:])
 
+        if secret == 'denied':
+            # Mother says no
+            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
         try:
@@ -81,18 +105,18 @@ 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
+    # Check if the users exist or not
     if not session.get('id'):
         # Our current URL, to which mothership will redirect back including a sessionstart
         original_url = f'{request.host_url}{path}'
@@ -101,6 +125,8 @@ def all_routes(path):
         app.logger.info('Redirecting to mothership with %s', original_url)
         return redirect(f'{app.config["MOTHERSHIP"]}/login?redirect={original_url}&callback={callback_url}')
 
+    if path == '':
+        path = 'index.html'
     file_path = os.path.join(app.config['SERVE_DIR'], path)
     if os.path.isfile(file_path):
         app.logger.info('Serving file %s', str(file_path))
@@ -108,9 +134,4 @@ 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'
-
-
-@app.route('/')
-def index():
-   return redirect('/index.html')
+        return handle_error_code(404)