3 files changed, 56 insertions, 17 deletions

diff --git a/src/services/api/graphql/libs/__init__.py b/src/services/api/graphql/libs/__init__.py
new file mode 100644
index 000000000..e69de29bb
--- /dev/null
+++ b/src/services/api/graphql/libs/__init__.py
diff --git a/src/services/api/graphql/libs/key_auth.py b/src/services/api/graphql/libs/key_auth.py
index 2db0f7d48..ffd7f32b2 100644
--- a/src/services/api/graphql/libs/key_auth.py
+++ b/src/services/api/graphql/libs/key_auth.py
@@ -1,5 +1,21 @@
+# 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/>.
+
+
+from ...session import SessionState
 
-from .. import state
 
 def check_auth(key_list, key):
     if not key_list:
@@ -10,9 +26,11 @@ def check_auth(key_list, key):
             key_id = k['id']
     return key_id
 
+
 def auth_required(key):
+    state = SessionState()
     api_keys = None
-    api_keys = state.settings['app'].state.vyos_keys
+    api_keys = state.keys
     key_id = check_auth(api_keys, key)
-    state.settings['app'].state.vyos_id = key_id
+    state.id = key_id
     return key_id
diff --git a/src/services/api/graphql/libs/token_auth.py b/src/services/api/graphql/libs/token_auth.py
index 8585485c9..4f743a096 100644
--- a/src/services/api/graphql/libs/token_auth.py
+++ b/src/services/api/graphql/libs/token_auth.py
@@ -1,46 +1,67 @@
+# 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")
+            return {'errors': ['missing secret']}
+        token = jwt.encode(payload=payload_data, key=secret, algorithm='HS256')
 
         users |= {user_id: user}
         return {'token': token}
     else:
-        return {"errors": ['failed pam authentication']}
+        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,8 +69,8 @@ def get_user_context(request):
             return context
 
         try:
-            secret = state.settings.get('secret')
-            payload = jwt.decode(token, secret, algorithms=["HS256"])
+            secret = getattr(state, 'secret', None)
+            payload = jwt.decode(token, secret, algorithms=['HS256'])
             user_id: str = payload.get('sub')
             if user_id is None:
                 return context
@@ -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