2 files changed, 70 insertions, 8 deletions

diff --git a/gn_auth/auth/authentication/oauth2/models/jwt_bearer_token.py b/gn_auth/auth/authentication/oauth2/models/jwt_bearer_token.py
index 2606ac6..71769e1 100644
--- a/gn_auth/auth/authentication/oauth2/models/jwt_bearer_token.py
+++ b/gn_auth/auth/authentication/oauth2/models/jwt_bearer_token.py
@@ -1,15 +1,50 @@
 """Implement model for JWTBearerToken"""
 import uuid
+import time
+from typing import Optional
 
 from authlib.oauth2.rfc7523 import JWTBearerToken as _JWTBearerToken
 
 from gn_auth.auth.db.sqlite3 import with_db_connection
 from gn_auth.auth.authentication.users import user_by_id
+from gn_auth.auth.authentication.oauth2.models.oauth2client import (
+    client as fetch_client)
 
 class JWTBearerToken(_JWTBearerToken):
     """Overrides default JWTBearerToken class."""
 
     def __init__(self, payload, header, options=None, params=None):
+        """Initialise the bearer token."""
+        # TOD0: Maybe remove this init and make this a dataclass like the way
+        #       OAuth2Client is a dataclass
         super().__init__(payload, header, options, params)
         self.user = with_db_connection(
             lambda conn:user_by_id(conn, uuid.UUID(payload["sub"])))
+        self.client = with_db_connection(
+            lambda conn: fetch_client(
+                conn, uuid.UUID(payload["oauth2_client_id"])
+            )
+        ).maybe(None, lambda _client: _client)
+
+
+    def check_client(self, client):
+        """Check that the client is right."""
+        return self.client.get_client_id() == client.get_client_id()
+
+
+    def get_expires_in(self) -> Optional[int]:
+        """Return the number of seconds the token is valid for since issue.
+
+        If `None`, the token never expires."""
+        if "exp" in self:
+            return self['exp'] - self['iat']
+        return None
+
+
+    def is_expired(self):
+        """Check whether the token is expired.
+
+        If there is no 'exp' member, assume this token will never expire."""
+        if "exp" in self:
+            return self["exp"] < time.time()
+        return False
diff --git a/gn_auth/auth/authentication/oauth2/models/oauth2client.py b/gn_auth/auth/authentication/oauth2/models/oauth2client.py
index 8fac648..1639e2e 100644
--- a/gn_auth/auth/authentication/oauth2/models/oauth2client.py
+++ b/gn_auth/auth/authentication/oauth2/models/oauth2client.py
@@ -1,18 +1,19 @@
 """OAuth2 Client model."""
 import json
-import logging
 import datetime
 from uuid import UUID
-from dataclasses import dataclass
 from functools import cached_property
-from typing import Sequence, Optional
+from dataclasses import asdict, dataclass
+from typing import Any, Sequence, Optional
 
 import requests
+from flask import current_app as app
 from requests.exceptions import JSONDecodeError
 from authlib.jose import KeySet, JsonWebKey
 from authlib.oauth2.rfc6749 import ClientMixin
 from pymonad.maybe import Just, Maybe, Nothing
 
+from gn_auth.debug import __pk__
 from gn_auth.auth.db import sqlite3 as db
 from gn_auth.auth.errors import NotFoundError
 from gn_auth.auth.authentication.users import (User,
@@ -62,23 +63,27 @@ class OAuth2Client(ClientMixin):
     def jwks(self) -> KeySet:
         """Return this client's KeySet."""
         jwksuri = self.client_metadata.get("public-jwks-uri")
+        __pk__(f"PUBLIC JWKs link for client {self.client_id}", jwksuri)
         if not bool(jwksuri):
-            logging.debug("No Public JWKs URI set for client!")
+            app.logger.debug("No Public JWKs URI set for client!")
             return KeySet([])
         try:
             ## IMPORTANT: This can cause a deadlock if the client is working in
             ##            single-threaded mode, i.e. can only serve one request
             ##            at a time.
             return KeySet([JsonWebKey.import_key(key)
-                           for key in requests.get(jwksuri).json()["jwks"]])
+                           for key in requests.get(
+                                   jwksuri,
+                                   timeout=300,
+                                   allow_redirects=True).json()["jwks"]])
         except requests.ConnectionError as _connerr:
-            logging.debug(
+            app.logger.debug(
                 "Could not connect to provided URI: %s", jwksuri, exc_info=True)
         except JSONDecodeError as _jsonerr:
-            logging.debug(
+            app.logger.debug(
                 "Could not convert response to JSON", exc_info=True)
         except Exception as _exc:# pylint: disable=[broad-except]
-            logging.debug(
+            app.logger.debug(
                 "Error retrieving the JWKs for the client.", exc_info=True)
         return KeySet([])
 
@@ -289,3 +294,25 @@ def delete_client(
         cursor.execute("DELETE FROM oauth2_tokens WHERE client_id=?", params)
         cursor.execute("DELETE FROM oauth2_clients WHERE client_id=?", params)
         return the_client
+
+
+def update_client_attribute(
+        client: OAuth2Client,# pylint: disable=[redefined-outer-name]
+        attribute: str,
+        value: Any
+) -> OAuth2Client:
+    """Return a new OAuth2Client with the given attribute updated/changed."""
+    attrs = {
+        attr: type(value)
+        for attr, value in asdict(client).items()
+        if attr != "client_id"
+    }
+    assert (
+        attribute in attrs.keys() and isinstance(value, attrs[attribute])), (
+            "Invalid attribute/value provided!")
+    return OAuth2Client(
+        client_id=client.client_id,
+        **{
+            attr: (value if attr==attribute else getattr(client, attr))
+            for attr in attrs
+        })