update focus

id.saladeaula.digital/app/integrations/apigateway_oauth2/authorization_server.py

@@ -28,6 +28,12 @@ GRANT_TYPES_EXPIRES_IN = {
 logger = Logger(__name__)
 
 
+def expires_in(client, grant_type: str) -> int:
+    return GRANT_TYPES_EXPIRES_IN.get(
+        grant_type, BearerTokenGenerator.DEFAULT_EXPIRES_IN
+    )
+
+
 class AuthorizationServer(oauth2.AuthorizationServer):
     def __init__(
         self,
@@ -36,15 +42,9 @@ class AuthorizationServer(oauth2.AuthorizationServer):
     ) -> None:
         self._persistence_layer = persistence_layer
 
-        super().__init__(
-            scopes_supported=(
-                set(OAUTH2_SCOPES_SUPPORTED.split()) if OAUTH2_SCOPES_SUPPORTED else []
-            )
-        )
-
-        def expires_in(client, grant_type: str) -> int:
-            return GRANT_TYPES_EXPIRES_IN.get(
-                grant_type, BearerTokenGenerator.DEFAULT_EXPIRES_IN
+        if OAUTH2_SCOPES_SUPPORTED:
+            super().__init__(
+                scopes_supported=set(OAUTH2_SCOPES_SUPPORTED.split()),
             )
 
         self.register_token_generator(
@@ -61,6 +61,7 @@ class AuthorizationServer(oauth2.AuthorizationServer):
         token: dict,
         request: requests.OAuth2Request,
     ) -> None:
+        """Define function to save the generated token into database."""
         if not request.payload:
             raise ValueError('Missing request payload')
 
@@ -88,6 +89,7 @@ class AuthorizationServer(oauth2.AuthorizationServer):
                 item={
                     'id': user_id,
                     'sk': f'SESSION#ACCESS_TOKEN#{access_token}',
+                    'client_id': client_id,
                     'ttl': access_token_ttl,
                     'created_at': now_,
                 }
@@ -112,6 +114,7 @@ class AuthorizationServer(oauth2.AuthorizationServer):
                     item={
                         'id': user_id,
                         'sk': f'SESSION#REFRESH_TOKEN#{refresh_token}',
+                        'client_id': client_id,
                         'ttl': access_token_ttl,
                         'created_at': now_,
                     }
@@ -133,7 +136,11 @@ class AuthorizationServer(oauth2.AuthorizationServer):
 
         return None
 
-    def query_client(self, client_id: str):
+    def query_client(
+        self,
+        client_id: str,
+    ):
+        """Query OAuth client by client_id."""
         client = self._persistence_layer.collection.get_item(
             KeyPair(
                 pk='OAUTH2',

id.saladeaula.digital/app/integrations/apigateway_oauth2/client.py

@@ -30,7 +30,10 @@ class OAuth2Client(ClientMixin):
     def get_client_id(self):
         return self.client_id
 
-    def get_allowed_scope(self, scope: Collection[str]) -> str:
+    def get_allowed_scope(
+        self,
+        scope: Collection[str],
+    ) -> str:
         if not scope:
             return ''
 
@@ -44,20 +47,36 @@ class OAuth2Client(ClientMixin):
 
         raise ValueError('Missing redirect_uris')
 
-    def check_response_type(self, response_type: str) -> bool:
+    def check_response_type(
+        self,
+        response_type: str,
+    ) -> bool:
         return response_type in self.response_types
 
-    def check_redirect_uri(self, redirect_uri: str) -> bool:
+    def check_redirect_uri(
+        self,
+        redirect_uri: str,
+    ) -> bool:
         return redirect_uri in self.redirect_uris
 
-    def check_endpoint_auth_method(self, method: str, endpoint: str) -> bool:
+    def check_endpoint_auth_method(
+        self,
+        method: str,
+        endpoint: str,
+    ) -> bool:
         if endpoint == 'token':
             return self.token_endpoint_auth_method == method
 
         return True
 
-    def check_grant_type(self, grant_type: str) -> bool:
+    def check_grant_type(
+        self,
+        grant_type: str,
+    ) -> bool:
         return grant_type in self.grant_types
 
-    def check_client_secret(self, client_secret: str) -> bool:
+    def check_client_secret(
+        self,
+        client_secret: str,
+    ) -> bool:
         return secrets.compare_digest(self.client_secret, client_secret)