perf: 优化user secret key 处理逻辑

apps/audits/api.py

@@ -1,6 +1,5 @@
 # -*- coding: utf-8 -*-
 #
-
 from importlib import import_module
 
 from django.conf import settings

apps/authentication/api/__init__.py

@@ -2,13 +2,14 @@
 #
 
 from .access_key import *
+from .common import *
 from .confirm import *
 from .connection_token import *
 from .feishu import *
 from .login_confirm import *
 from .mfa import *
 from .password import *
+from .session import *
 from .sso import *
 from .temp_token import *
 from .token import *
-from .common import *

apps/authentication/api/login_confirm.py

@@ -9,6 +9,7 @@ from common.utils import get_logger
 from .. import errors, mixins
 
 __all__ = ['TicketStatusApi']
+
 logger = get_logger(__name__)
 
 

apps/authentication/api/session.py

@@ -0,0 +1,62 @@
+import time
+from threading import Thread
+
+from django.conf import settings
+from django.contrib.auth import logout
+from rest_framework import generics
+from rest_framework import status
+from rest_framework.response import Response
+
+from common.permissions import IsValidUser
+from common.sessions.cache import user_session_manager
+from common.utils import get_logger
+
+__all__ = ['UserSessionApi']
+
+logger = get_logger(__name__)
+
+
+class UserSessionManager:
+
+    def __init__(self, request):
+        self.request = request
+        self.session = request.session
+
+    def connect(self):
+        user_session_manager.add_or_increment(self.session.session_key)
+
+    def disconnect(self):
+        user_session_manager.decrement_or_remove(self.session.session_key)
+        if self.should_delete_session():
+            thread = Thread(target=self.delay_delete_session)
+            thread.start()
+
+    def should_delete_session(self):
+        return (self.session.modified or settings.SESSION_SAVE_EVERY_REQUEST) and \
+            not self.session.is_empty() and \
+            self.session.get_expire_at_browser_close() and \
+            not user_session_manager.check_active(self.session.session_key)
+
+    def delay_delete_session(self):
+        timeout = 6
+        check_interval = 0.5
+
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            time.sleep(check_interval)
+            if user_session_manager.check_active(self.session.session_key):
+                return
+
+        logout(self.request)
+
+
+class UserSessionApi(generics.RetrieveDestroyAPIView):
+    permission_classes = (IsValidUser,)
+
+    def retrieve(self, request, *args, **kwargs):
+        UserSessionManager(request).connect()
+        return Response(status=status.HTTP_200_OK)
+
+    def destroy(self, request, *args, **kwargs):
+        UserSessionManager(request).disconnect()
+        return Response(status=status.HTTP_204_NO_CONTENT)

apps/authentication/urls/api_urls.py

@@ -32,6 +32,7 @@ urlpatterns = [
     path('password/reset-code/', api.UserResetPasswordSendCodeApi.as_view(), name='reset-password-code'),
     path('password/verify/', api.UserPasswordVerifyApi.as_view(), name='user-password-verify'),
     path('login-confirm-ticket/status/', api.TicketStatusApi.as_view(), name='login-confirm-ticket-status'),
+    path('user-session/', api.UserSessionApi.as_view(), name='user-session'),
 ]
 
 urlpatterns += router.urls + passkey_urlpatterns

apps/common/sessions/cache.py

@@ -1,16 +1,19 @@
 import re
+from importlib import import_module
 
+from django.conf import settings
 from django.contrib.sessions.backends.cache import (
     SessionStore as DjangoSessionStore
 )
-from django.core.cache import cache
+from django.core.cache import cache, caches
 
 from jumpserver.utils import get_current_request
 
 
 class SessionStore(DjangoSessionStore):
     ignore_urls = [
-        r'^/api/v1/users/profile/'
+        r'^/api/v1/users/profile/',
+        r'^/api/v1/authentication/user-session/'
     ]
 
     def __init__(self, *args, **kwargs):
@@ -55,12 +58,12 @@ class RedisUserSessionManager:
             session_keys.append(key)
         return session_keys
 
-    def get_keys(self):
-        session_keys = []
-        for k in self.client.hgetall(self.JMS_SESSION_KEY).keys():
-            key = k.decode('utf-8')
-            session_keys.append(key)
-        return session_keys
+    @staticmethod
+    def get_keys():
+        session_store_cls = import_module(settings.SESSION_ENGINE).SessionStore
+        cache_key_prefix = session_store_cls.cache_key_prefix
+        keys = caches[settings.SESSION_CACHE_ALIAS].iter_keys('*')
+        return [k.replace(cache_key_prefix, '') for k in keys]
 
 
 user_session_manager = RedisUserSessionManager()

apps/notifications/ws.py

@@ -1,12 +1,8 @@
 import json
-import time
-from threading import Thread
 
 from channels.generic.websocket import JsonWebsocketConsumer
-from django.conf import settings
 
 from common.db.utils import safe_db_connection
-from common.sessions.cache import user_session_manager
 from common.utils import get_logger
 from .signal_handlers import new_site_msg_chan
 from .site_msg import SiteMessageUtil
@@ -26,7 +22,6 @@ class SiteMsgWebsocket(JsonWebsocketConsumer):
         user = self.scope["user"]
         if user.is_authenticated:
             self.accept()
-            user_session_manager.add_or_increment(self.session.session_key)
             self.sub = self.watch_recv_new_site_msg()
         else:
             self.close()
@@ -70,32 +65,3 @@ class SiteMsgWebsocket(JsonWebsocketConsumer):
         if not self.sub:
             return
         self.sub.unsubscribe()
-
-        user_session_manager.decrement_or_remove(self.session.session_key)
-        if self.should_delete_session():
-            thread = Thread(target=self.delay_delete_session)
-            thread.start()
-
-    def should_delete_session(self):
-        return (self.session.modified or settings.SESSION_SAVE_EVERY_REQUEST) and \
-            not self.session.is_empty() and \
-            self.session.get_expire_at_browser_close() and \
-            not user_session_manager.check_active(self.session.session_key)
-
-    def delay_delete_session(self):
-        timeout = 6
-        check_interval = 0.5
-
-        start_time = time.time()
-        while time.time() - start_time < timeout:
-            time.sleep(check_interval)
-            if user_session_manager.check_active(self.session.session_key):
-                return
-
-        self.delete_session()
-
-    def delete_session(self):
-        try:
-            self.session.delete()
-        except Exception as e:
-            logger.info(f'delete session error: {e}')