From b023ca0c692759e83f0329f5f2ee7fbdceff220a Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Thu, 21 Sep 2023 14:02:09 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20saml=20=E7=94=A8=E6=88=B7=E6=B2=A1?=
 =?UTF-8?q?=E7=8E=B0=E5=9C=A8=E8=AE=B0=E5=BD=95=20(#11641)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: feng <1304903146@qq.com>
---
 apps/audits/api.py                       |  5 ++++-
 apps/audits/models.py                    | 10 ++++++++++
 apps/audits/signal_handlers/login_log.py |  9 +++------
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/apps/audits/api.py b/apps/audits/api.py
index e0d7dbdec..72c1d8a99 100644
--- a/apps/audits/api.py
+++ b/apps/audits/api.py
@@ -268,7 +268,10 @@ class UserSessionViewSet(CommonApiMixin, viewsets.ModelViewSet):
         return user_ids
 
     def get_queryset(self):
-        queryset = UserSession.objects.filter(date_expired__gt=timezone.now())
+        keys = UserSession.get_keys()
+        queryset = UserSession.objects.filter(
+            date_expired__gt=timezone.now(), key__in=keys
+        )
         if current_org.is_root():
             return queryset
         user_ids = self.org_user_ids
diff --git a/apps/audits/models.py b/apps/audits/models.py
index 0ec6f2809..34ec301f2 100644
--- a/apps/audits/models.py
+++ b/apps/audits/models.py
@@ -1,7 +1,9 @@
 import os
 import uuid
+from importlib import import_module
 
 from django.conf import settings
+from django.core.cache import caches
 from django.db import models
 from django.db.models import Q
 from django.utils import timezone
@@ -269,9 +271,17 @@ class UserSession(models.Model):
     def backend_display(self):
         return gettext(self.backend)
 
+    @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].keys('*')
+        return [k.replace(cache_key_prefix, '') for k in keys]
+
     @classmethod
     def clear_expired_sessions(cls):
         cls.objects.filter(date_expired__lt=timezone.now()).delete()
+        cls.objects.exclude(key__in=cls.get_keys()).delete()
 
     class Meta:
         ordering = ['-date_created']
diff --git a/apps/audits/signal_handlers/login_log.py b/apps/audits/signal_handlers/login_log.py
index 8ce8abff9..fae32a44b 100644
--- a/apps/audits/signal_handlers/login_log.py
+++ b/apps/audits/signal_handlers/login_log.py
@@ -5,7 +5,6 @@ from importlib import import_module
 
 from django.conf import settings
 from django.contrib.auth import BACKEND_SESSION_KEY
-from django.core.cache import caches
 from django.dispatch import receiver
 from django.utils import timezone, translation
 from django.utils.functional import LazyObject
@@ -83,11 +82,10 @@ def generate_data(username, request, login_type=None):
 
 
 def create_user_session(request, user_id, instance: UserLoginLog):
-    session_key = request.session.session_key
+    session_key = request.session.session_key or '-'
     session_store_cls = import_module(settings.SESSION_ENGINE).SessionStore
     session_store = session_store_cls(session_key=session_key)
-    cache_key = session_store.cache_key
-    ttl = caches[settings.SESSION_CACHE_ALIAS].ttl(cache_key)
+    ttl = session_store.get_expiry_age()
 
     online_session_data = {
         'user_id': user_id,
@@ -114,9 +112,8 @@ def on_user_auth_success(sender, user, request, login_type=None, **kwargs):
     request.session['login_time'] = data['datetime'].strftime("%Y-%m-%d %H:%M:%S")
     data.update({'mfa': int(user.mfa_enabled), 'status': True})
     instance = write_login_log(**data)
-    session_key = request.session.session_key
     # TODO 目前只记录 web 登录的 session
-    if not session_key or instance.type != LoginTypeChoices.web:
+    if instance.type != LoginTypeChoices.web:
         return
     create_user_session(request, user.id, instance)