From 50cf3dd2318718c995f9612abba4de448f9fb942 Mon Sep 17 00:00:00 2001
From: EC2 Default User <ec2-user@ip-172-31-30-15.us-west-2.compute.internal>
Date: Tue, 24 Dec 2024 00:38:21 +0000
Subject: [PATCH] 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             |  4 ++++
 policyengine_api/api.py          | 10 +++++++++
 policyengine_api/auth_context.py | 37 ++++++++++++++++++++++++++++++++
 policyengine_api/validator.py    | 24 +++++++++++++++++++++
 setup.py                         |  1 +
 5 files changed, 76 insertions(+)
 create mode 100644 policyengine_api/auth_context.py
 create mode 100644 policyengine_api/validator.py

diff --git a/changelog_entry.yaml b/changelog_entry.yaml
index e69de29b..8bf9625b 100644
--- a/changelog_entry.yaml
+++ b/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.
diff --git a/policyengine_api/api.py b/policyengine_api/api.py
index 684bb3f9..4386adb0 100644
--- a/policyengine_api/api.py
+++ b/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",
diff --git a/policyengine_api/auth_context.py b/policyengine_api/auth_context.py
new file mode 100644
index 00000000..f8d758b7
--- /dev/null
+++ b/policyengine_api/auth_context.py
@@ -0,0 +1,37 @@
+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
diff --git a/policyengine_api/validator.py b/policyengine_api/validator.py
new file mode 100644
index 00000000..74a2e127
--- /dev/null
+++ b/policyengine_api/validator.py
@@ -0,0 +1,24 @@
+# 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},
+        }
diff --git a/setup.py b/setup.py
index c28b9bcb..659468d2 100644
--- a/setup.py
+++ b/setup.py
@@ -32,6 +32,7 @@
         "streamlit",
         "werkzeug",
         "Flask-Caching>=2,<3",
+        "Authlib",
     ],
     extras_require={
         "dev": [