1 files changed, 31 insertions, 10 deletions
diff --git a/src/services/api/graphql/libs/token_auth.py b/src/services/api/graphql/libs/token_auth.py
index 8585485c9..2d772e035 100644
--- a/src/services/api/graphql/libs/token_auth.py
+++ b/src/services/api/graphql/libs/token_auth.py
@@ -1,33 +1,52 @@
+# Copyright 2021-2024 VyOS maintainers and contributors <maintainers@vyos.io>
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License as published by the Free Software Foundation; either
+# version 2.1 of the License, or (at your option) any later version.
+#
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this library.  If not, see <http://www.gnu.org/licenses/>.
+
+
 import jwt
 import uuid
 import pam
 from secrets import token_hex
 
-from .. import state
+from ... session import SessionState
+
 
 def _check_passwd_pam(username: str, passwd: str) -> bool:
     if pam.authenticate(username, passwd):
         return True
     return False
 
+
 def init_secret():
-    length = int(state.settings['app'].state.vyos_secret_len)
+    state = SessionState()
+    length = int(state.secret_len)
     secret = token_hex(length)
-    state.settings['secret'] = secret
+    state.secret = secret
+
 
 def generate_token(user: str, passwd: str, secret: str, exp: int) -> dict:
     if user is None or passwd is None:
         return {}
+    state = SessionState()
     if _check_passwd_pam(user, passwd):
-        app = state.settings['app']
         try:
-            users = app.state.vyos_token_users
+            users = state.token_users
         except AttributeError:
-            app.state.vyos_token_users = {}
-            users = app.state.vyos_token_users
+            users = state.token_users = {}
         user_id = uuid.uuid1().hex
         payload_data = {'iss': user, 'sub': user_id, 'exp': exp}
-        secret = state.settings.get('secret')
+        secret = getattr(state, 'secret', None)
         if secret is None:
             return {"errors": ['missing secret']}
         token = jwt.encode(payload=payload_data, key=secret, algorithm="HS256")
@@ -37,10 +56,12 @@ def generate_token(user: str, passwd: str, secret: str, exp: int) -> dict:
     else:
         return {"errors": ['failed pam authentication']}
 
+
 def get_user_context(request):
     context = {}
     context['request'] = request
     context['user'] = None
+    state = SessionState()
     if 'Authorization' in request.headers:
         auth = request.headers['Authorization']
         scheme, token = auth.split()
@@ -48,7 +69,7 @@ def get_user_context(request):
             return context
 
         try:
-            secret = state.settings.get('secret')
+            secret = getattr(state, 'secret', None)
             payload = jwt.decode(token, secret, algorithms=["HS256"])
             user_id: str = payload.get('sub')
             if user_id is None:
@@ -59,7 +80,7 @@ def get_user_context(request):
         except jwt.PyJWTError:
             return context
         try:
-            users = state.settings['app'].state.vyos_token_users
+            users = state.token_users
         except AttributeError:
             return context