From 24ed57b98a84114185d35cc0963b422218874d2b Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Mon, 18 Sep 2023 16:20:55 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=B8=89=E6=96=B9=E7=99=BB=E5=BD=95?=
 =?UTF-8?q?=E7=94=A8=E6=88=B7=E6=97=A0=E6=B3=95=E4=B8=8B=E7=BA=BF=20(#1159?=
 =?UTF-8?q?2)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: feng <1304903146@qq.com>
---
 apps/audits/signal_handlers/login_log.py | 8 +++++---
 apps/authentication/signal_handlers.py   | 5 +++++
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/apps/audits/signal_handlers/login_log.py b/apps/audits/signal_handlers/login_log.py
index d1bd6b58b..8ce8abff9 100644
--- a/apps/audits/signal_handlers/login_log.py
+++ b/apps/audits/signal_handlers/login_log.py
@@ -82,7 +82,8 @@ def generate_data(username, request, login_type=None):
     return data
 
 
-def create_user_session(session_key, user_id, instance: UserLoginLog):
+def create_user_session(request, user_id, instance: UserLoginLog):
+    session_key = request.session.session_key
     session_store_cls = import_module(settings.SESSION_ENGINE).SessionStore
     session_store = session_store_cls(session_key=session_key)
     cache_key = session_store.cache_key
@@ -99,7 +100,8 @@ def create_user_session(session_key, user_id, instance: UserLoginLog):
         'date_created': instance.datetime,
         'date_expired': instance.datetime + timedelta(seconds=ttl),
     }
-    UserSession.objects.create(**online_session_data)
+    user_session = UserSession.objects.create(**online_session_data)
+    request.session['user_session_id'] = user_session.id
 
 
 @receiver(post_auth_success)
@@ -116,7 +118,7 @@ def on_user_auth_success(sender, user, request, login_type=None, **kwargs):
     # TODO 目前只记录 web 登录的 session
     if not session_key or instance.type != LoginTypeChoices.web:
         return
-    create_user_session(session_key, user.id, instance)
+    create_user_session(request, user.id, instance)
 
 
 @receiver(post_auth_failed)
diff --git a/apps/authentication/signal_handlers.py b/apps/authentication/signal_handlers.py
index b1f55f689..3ac92411f 100644
--- a/apps/authentication/signal_handlers.py
+++ b/apps/authentication/signal_handlers.py
@@ -7,6 +7,7 @@ from django.dispatch import receiver
 from django_cas_ng.signals import cas_user_authenticated
 
 from apps.jumpserver.settings.auth import AUTHENTICATION_BACKENDS_THIRD_PARTY
+from audits.models import UserSession
 from .signals import post_auth_success, post_auth_failed, user_auth_failed, user_auth_success
 
 
@@ -23,6 +24,9 @@ def on_user_auth_login_success(sender, user, request, **kwargs):
     if not request.session.get("auth_third_party_done") and \
             request.session.get('auth_backend') in AUTHENTICATION_BACKENDS_THIRD_PARTY:
         request.session['auth_third_party_required'] = 1
+
+    user_session_id = request.session.get('user_session_id')
+    UserSession.objects.filter(id=user_session_id).update(key=request.session.session_key)
     # 单点登录，超过了自动退出
     if settings.USER_LOGIN_SINGLE_MACHINE_ENABLED:
         lock_key = 'single_machine_login_' + str(user.id)
@@ -30,6 +34,7 @@ def on_user_auth_login_success(sender, user, request, **kwargs):
         if session_key and session_key != request.session.session_key:
             session = import_module(settings.SESSION_ENGINE).SessionStore(session_key)
             session.delete()
+            UserSession.objects.filter(key=session_key).delete()
         cache.set(lock_key, request.session.session_key, None)
 
     # 标记登录，设置 cookie，前端可以控制刷新, Middleware 会拦截这个生成 cookie