From 0a28c5650cae8c75fb8b8af676044d2a68744e60 Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Mon, 16 Oct 2023 11:28:53 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E4=B8=89=E6=96=B9=E7=94=A8=E6=88=B7?=
 =?UTF-8?q?=E7=99=BB=E5=BD=95=E9=80=9A=E7=9F=A5=20(#11846)?=
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 | 32 ++++++++++++++----------
 apps/authentication/middleware.py        |  7 ++++++
 apps/authentication/mixins.py            |  3 ++-
 3 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/apps/audits/signal_handlers/login_log.py b/apps/audits/signal_handlers/login_log.py
index 835e4842f..2ad9c0a7a 100644
--- a/apps/audits/signal_handlers/login_log.py
+++ b/apps/audits/signal_handlers/login_log.py
@@ -11,7 +11,6 @@ from django.utils.functional import LazyObject
 from django.utils.translation import gettext_lazy as _
 from rest_framework.request import Request
 
-from acls.const import ActionChoices
 from acls.models import LoginACL
 from acls.notifications import UserLoginReminderMsg
 from audits.models import UserLoginLog
@@ -85,6 +84,9 @@ def generate_data(username, request, login_type=None):
 
 
 def create_user_session(request, user_id, instance: UserLoginLog):
+    # TODO 目前只记录 web 登录的 session
+    if instance.type != LoginTypeChoices.web:
+        return
     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)
@@ -102,10 +104,21 @@ def create_user_session(request, user_id, instance: UserLoginLog):
         'date_expired': instance.datetime + timedelta(seconds=ttl),
     }
     user_session = UserSession.objects.create(**online_session_data)
-    request.session['user_session_id'] = user_session.id
+    request.session['user_session_id'] = str(user_session.id)
 
 
-def send_login_info_to_reviewers(instance: UserLoginLog, reviewers):
+def send_login_info_to_reviewers(instance: UserLoginLog | str, auth_acl_id):
+    if isinstance(instance, str):
+        instance = UserLoginLog.objects.filter(id=instance).first()
+
+    if not instance:
+        return
+
+    acl = LoginACL.objects.filter(id=auth_acl_id).first()
+    if not acl or not acl.reviewers.exists():
+        return
+
+    reviewers = acl.reviewers.all()
     for reviewer in reviewers:
         UserLoginReminderMsg(reviewer, instance).publish_async()
 
@@ -119,22 +132,15 @@ def on_user_auth_success(sender, user, request, login_type=None, **kwargs):
     data.update({'mfa': int(user.mfa_enabled), 'status': True})
     instance = write_login_log(**data)
 
-    # TODO 目前只记录 web 登录的 session
-    if instance.type != LoginTypeChoices.web:
-        return
     create_user_session(request, user.id, instance)
-
+    request.session['user_log_id'] = str(instance.id)
+    request.session['can_send_notifications'] = True
     auth_notice_required = request.session.get('auth_notice_required')
     if not auth_notice_required:
         return
 
     auth_acl_id = request.session.get('auth_acl_id')
-    acl = LoginACL.objects.filter(id=auth_acl_id, action=ActionChoices.notice).first()
-    if not acl or not acl.reviewers.exists():
-        return
-
-    reviewers = acl.reviewers.all()
-    send_login_info_to_reviewers(instance, reviewers)
+    send_login_info_to_reviewers(instance, auth_acl_id)
 
 
 @receiver(post_auth_failed)
diff --git a/apps/authentication/middleware.py b/apps/authentication/middleware.py
index 572ad39dc..9f245099d 100644
--- a/apps/authentication/middleware.py
+++ b/apps/authentication/middleware.py
@@ -8,6 +8,7 @@ from django.utils.deprecation import MiddlewareMixin
 from django.utils.translation import gettext as _
 
 from apps.authentication import mixins
+from audits.signal_handlers import send_login_info_to_reviewers
 from authentication.signals import post_auth_failed
 from common.utils import gen_key_pair
 from common.utils import get_request_ip
@@ -105,6 +106,12 @@ class ThirdPartyLoginMiddleware(mixins.AuthMixin):
                 guard_url = "%s?%s" % (guard_url, args)
             response = redirect(guard_url)
         finally:
+            if request.session.get('can_send_notifications') and \
+                    self.request.session['auth_notice_required']:
+                request.session['can_send_notifications'] = False
+                user_log_id = self.request.session.get('user_log_id')
+                auth_acl_id = self.request.session.get('auth_acl_id')
+                send_login_info_to_reviewers(user_log_id, auth_acl_id)
             return response
 
 
diff --git a/apps/authentication/mixins.py b/apps/authentication/mixins.py
index 2446b4667..31cb1dc19 100644
--- a/apps/authentication/mixins.py
+++ b/apps/authentication/mixins.py
@@ -523,7 +523,8 @@ class AuthMixin(CommonMixin, AuthPreCheckMixin, AuthACLMixin, MFAMixin, AuthPost
     def clear_auth_mark(self):
         keys = [
             'auth_password', 'user_id', 'auth_confirm_required',
-            'auth_notice_required', 'auth_ticket_id', 'auth_acl_id'
+            'auth_notice_required', 'auth_ticket_id', 'auth_acl_id',
+            'user_session_id', 'user_log_id', 'can_send_notifications'
         ]
         for k in keys:
             self.request.session.pop(k, '')