1 files changed, 89 insertions, 0 deletions
diff --git a/gn_auth/auth/authentication/oauth2/grants/jwt_bearer_grant.py b/gn_auth/auth/authentication/oauth2/grants/jwt_bearer_grant.py
new file mode 100644
index 0000000..492009c
--- /dev/null
+++ b/gn_auth/auth/authentication/oauth2/grants/jwt_bearer_grant.py
@@ -0,0 +1,89 @@
+"""JWT as Authorisation Grant"""
+import uuid
+from urllib.parse import urlparse
+from datetime import datetime, timedelta
+
+from flask import request, current_app as app
+
+from authlib.jose import jwt
+from authlib.jose.rfc7517.key_set import KeySet
+from authlib.oauth2.rfc7523.jwt_bearer import JWTBearerGrant as _JWTBearerGrant
+from authlib.oauth2.rfc7523.token import (
+    JWTBearerTokenGenerator as _JWTBearerTokenGenerator)
+
+from gn_auth.auth.authentication.users import user_by_id
+from gn_auth.auth.db.sqlite3 import connection, with_db_connection
+from gn_auth.auth.authentication.oauth2.models.oauth2client import client
+from gn_auth.auth.authentication.oauth2.grants.authorisation_code_grant import AuthorisationCodeGrant
+
+
+class JWTBearerTokenGenerator(_JWTBearerTokenGenerator):
+    """
+    A JSON Web Token formatted bearer token generator for jwt-bearer grant type.
+    """
+
+    DEFAULT_EXPIRES_IN = 300
+
+    def get_token_data(self, grant_type, client, expires_in=300, user=None, scope=None):
+        """Post process data to prevent JSON serialization problems."""
+        tokendata = super().get_token_data(
+            grant_type, client, expires_in, user, scope)
+        return {
+            **{
+                key: str(value) if key.endswith("_id") else value
+                for key, value in tokendata.items()
+            },
+            "sub": str(tokendata["sub"])}
+
+
+class JWTBearerGrant(_JWTBearerGrant, AuthorisationCodeGrant):
+    """Implement JWT as Authorisation Grant."""
+
+
+    def create_authorization_response(self, redirect_uri: str, grant_user):
+        resp = super().create_authorization_response(redirect_uri, grant_user)
+        headers = dict(resp[2])
+        location = urlparse(headers["Location"])
+        query = {
+            key.strip(): value.strip() for key, value in
+            (item.split("=") for  item in
+             (param.strip() for param in location.query.split("&")))}
+        parsed_redirect = urlparse(redirect_uri)
+        issued = datetime.now()
+        jwtkey = app.config["JWT_PRIVATE_KEY"]
+        jwttoken = jwt.encode(
+            {"alg": "RS256", "typ": "jwt", "kid": jwtkey.kid},
+            {
+                "iss": str(self.client.client_id),
+                "sub": str(grant_user.user_id),
+                "aud": f"{parsed_redirect.scheme}://{parsed_redirect.netloc}",
+                "exp": (issued + timedelta(minutes=5)),
+                "nbf": int(issued.timestamp()),
+                "iat": int(issued.timestamp()),
+                "jti": str(uuid.uuid4()),
+                "code": query["code"]},
+            jwtkey).decode("utf8")
+        return (302, "", [("Location", f"{location.geturl()}&jwt={jwttoken}")])
+
+
+    def resolve_issuer_client(self, issuer):
+        """Fetch client via "iss" in assertion claims."""
+        return with_db_connection(
+            lambda conn: self.server.query_client(issuer))
+
+
+    def resolve_client_key(self, client, headers, payload):
+        """Resolve client key to decode assertion data."""
+        return app.config["JWT_PUBLIC_KEY"]
+
+
+    def authenticate_user(self, subject):
+        """Authenticate user with the given assertion claims."""
+        return with_db_connection(lambda conn: user_by_id(conn, subject))
+
+
+    def has_granted_permission(self, client, user):
+        """
+        Check if the client has permission to access the given user's resource.
+        """
+        return True # TODO: Check this!!!