Integrate JWT Auth into API

partially fixes #2063

This will allow us to add permissions checks to our API calls. Currently
it does not require a valid JWT, but if one exists will store the "sub"
field such that the user's id can be checked against our user table to
establish permissions.

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    added:
+      - API now attempts to parse a the bearer token if one is provided and logs success/fail.

policyengine_api/api.py

@@ -7,8 +7,11 @@
 import flask
 import yaml
 from flask_caching import Cache
+from authlib.integrations.flask_oauth2 import ResourceProtector
+from policyengine_api.validator import Auth0JWTBearerTokenValidator
 from policyengine_api.utils import make_cache_key
 from .constants import VERSION
+import policyengine_api.auth_context as auth_context
 
 # from werkzeug.middleware.profiler import ProfilerMiddleware
 
@@ -42,6 +45,13 @@
 
 app = application = flask.Flask(__name__)
 
+## as per https://auth0.com/docs/quickstart/backend/python/interactive
+require_auth = ResourceProtector()
+validator = Auth0JWTBearerTokenValidator()
+require_auth.register_token_validator(validator)
+
+auth_context.configure(app, require_auth=require_auth)
+
 app.config.from_mapping(
     {
         "CACHE_TYPE": "RedisCache",

policyengine_api/auth_context.py

@@ -0,0 +1,41 @@
+from flask import Flask, g
+from werkzeug.local import LocalProxy
+from authlib.integrations.flask_oauth2 import ResourceProtector
+
+
+def configure(app: Flask, require_auth: ResourceProtector):
+    """
+    Configure the application to attempt to get and validate a bearer token.
+    If there is a token and it's valid the user id is added to the request context
+    which can be accessed via get_user_id
+    Otherwise, the request is accepted but get_user_id returns None
+
+    This supports our current auth model where only user-specific actions are restricted and
+    then only to allow the user
+    """
+
+    # If the user is authenticated then get the user id from the token
+    # And add it to the flask request context.
+    @app.before_request
+    def get_user():
+        try:
+            token = require_auth.acquire_token()
+            print(f"Validated JWT for sub {g.authlib_server_oauth2_token.sub}")
+        except Exception as ex:
+            print(f"Unable to parse a valid bearer token from request: {ex}")
+
+
+def get_user() -> None | str:
+    # I didn't see this documented anywhere, but if you look at the source code
+    # the validator stores the token in the flask global context under this name.
+    if "authlib_server_oauth2_token" not in g:
+        print(
+            "authlib_server_oauth2_token is not in the flask global context. Please make sure you called 'configure' on the app"
+        )
+        return None
+    if "sub" not in g.authlib_server_oauth2_token:
+        print(
+            "ERROR: authlib_server_oauth2_token does not contain a sub field. The JWT validator should force this to be true."
+        )
+        return None
+    return g.authlib_server_oauth2_token.sub

policyengine_api/validator.py

@@ -0,0 +1,26 @@
+# As defined by https://auth0.com/docs/quickstart/backend/python/interactive
+import json
+from urllib.request import urlopen
+
+from authlib.oauth2.rfc7523 import JWTBearerTokenValidator
+from authlib.jose.rfc7517.jwk import JsonWebKey
+
+
+class Auth0JWTBearerTokenValidator(JWTBearerTokenValidator):
+    def __init__(
+        self,
+        audience="https://api.policyengine.org/",
+    ):
+        issuer = "https://policyengine.uk.auth0.com/"
+        jsonurl = urlopen(
+            f"https://policyengine.uk.auth0.com/.well-known/jwks.json"
+        )
+        public_key = JsonWebKey.import_key_set(json.loads(jsonurl.read()))
+        super(Auth0JWTBearerTokenValidator, self).__init__(public_key)
+        self.claims_options = {
+            "exp": {"essential": True},
+            "aud": {"essential": True, "value": audience},
+            "iss": {"essential": True, "value": issuer},
+            # Provides the user id as we currently use it.
+            "sub": {"essential": True},
+        }

setup.py

@@ -32,6 +32,7 @@
         "streamlit",
         "werkzeug",
         "Flask-Caching>=2,<3",
+        "Authlib",
     ],
     extras_require={
         "dev": [