From 0778a398941e60c1df4b72ea2e06c5b2d82f3f64 Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Thu, 16 Nov 2023 14:41:35 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E5=9C=A8=E7=BA=BF=E4=BC=9A=E8=AF=9D?=
 =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=B4=BB=E8=B7=83=E7=8A=B6=E6=80=81=E8=BF=87?=
 =?UTF-8?q?=E6=BB=A4=20(#12134)?=
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     |  3 ++-
 apps/audits/filters.py | 45 +++++++++++++++++++++---------------------
 2 files changed, 25 insertions(+), 23 deletions(-)

diff --git a/apps/audits/api.py b/apps/audits/api.py
index 35315aa65..c2cd84756 100644
--- a/apps/audits/api.py
+++ b/apps/audits/api.py
@@ -29,6 +29,7 @@ from terminal.models import default_storage
 from users.models import User
 from .backends import TYPE_ENGINE_MAPPING
 from .const import ActivityChoices
+from .filters import UserSessionFilterSet
 from .models import (
     FTPLog, UserLoginLog, OperateLog, PasswordChangeLog,
     ActivityLog, JobLog, UserSession
@@ -255,7 +256,7 @@ class PasswordChangeLogViewSet(OrgReadonlyModelViewSet):
 class UserSessionViewSet(CommonApiMixin, viewsets.ModelViewSet):
     http_method_names = ('get', 'post', 'head', 'options', 'trace')
     serializer_class = UserSessionSerializer
-    filterset_fields = ['id', 'ip', 'city', 'type']
+    filterset_class = UserSessionFilterSet
     search_fields = ['id', 'ip', 'city']
     rbac_perms = {
         'offline': ['audits.offline_usersession']
diff --git a/apps/audits/filters.py b/apps/audits/filters.py
index c15c22b56..d24cf72fd 100644
--- a/apps/audits/filters.py
+++ b/apps/audits/filters.py
@@ -1,11 +1,12 @@
-from django.db.models import F, Value
-from django.db.models.functions import Concat
-from django_filters.rest_framework import CharFilter
+from django.core.cache import cache
+from django_filters import rest_framework as drf_filters
 from rest_framework import filters
 from rest_framework.compat import coreapi, coreschema
 
-from orgs.utils import current_org
 from common.drf.filters import BaseFilterSet
+from notifications.ws import WS_SESSION_KEY
+from orgs.utils import current_org
+from .models import UserSession
 
 __all__ = ['CurrentOrgMembersFilter']
 
@@ -34,21 +35,21 @@ class CurrentOrgMembersFilter(filters.BaseFilterBackend):
             queryset = queryset.filter(user__in=self._get_user_list())
         return queryset
 
-#
-# class CommandExecutionFilter(BaseFilterSet):
-#     hostname_ip = CharFilter(method='filter_hostname_ip')
-#
-#     class Meta:
-#         model = CommandExecution.hosts.through
-#         fields = (
-#             'id', 'asset', 'commandexecution', 'hostname_ip'
-#         )
-#
-#     def filter_hostname_ip(self, queryset, name, value):
-#         queryset = queryset.annotate(
-#             hostname_ip=Concat(
-#                 F('asset__hostname'), Value('('),
-#                 F('asset__address'), Value(')')
-#             )
-#         ).filter(hostname_ip__icontains=value)
-#         return queryset
+
+class UserSessionFilterSet(BaseFilterSet):
+    is_active = drf_filters.BooleanFilter(method='filter_is_active')
+
+    @staticmethod
+    def filter_is_active(queryset, name, is_active):
+        redis_client = cache.client.get_client()
+        members = redis_client.smembers(WS_SESSION_KEY)
+        members = [member.decode('utf-8') for member in members]
+        if is_active:
+            queryset = queryset.filter(key__in=members)
+        else:
+            queryset = queryset.exclude(key__in=members)
+        return queryset
+
+    class Meta:
+        model = UserSession
+        fields = ['id', 'ip', 'city', 'type']