From 2cad97065f78958ee3c91dc0108f75b0670e2c53 Mon Sep 17 00:00:00 2001
From: wangruidong <940853815@qq.com>
Date: Thu, 30 Nov 2023 11:14:07 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E8=B5=84=E4=BA=A7=E8=AF=A6=E6=83=85?=
 =?UTF-8?q?=E9=A1=B5=E9=9D=A2=E6=B7=BB=E5=8A=A0=E5=8E=86=E5=8F=B2=E6=89=A7?=
 =?UTF-8?q?=E8=A1=8C=E5=91=BD=E4=BB=A4=E5=88=97=E8=A1=A8=E9=A1=B5=E9=9D=A2?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/terminal/api/session/command.py | 5 +++++
 apps/terminal/api/session/session.py | 2 +-
 apps/terminal/filters.py             | 3 ++-
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/apps/terminal/api/session/command.py b/apps/terminal/api/session/command.py
index 494d001c0..387aebc3e 100644
--- a/apps/terminal/api/session/command.py
+++ b/apps/terminal/api/session/command.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 #
+import uuid
 from django.utils import timezone
 from rest_framework import generics
 from rest_framework.fields import DateTimeField
@@ -166,6 +167,7 @@ class CommandViewSet(JMSBulkModelViewSet):
 
     def get_queryset(self):
         command_storage_id = self.request.query_params.get('command_storage_id')
+        asset_id = self.request.query_params.get('asset_id')
         if not command_storage_id:
             return Command.objects.none()
 
@@ -174,6 +176,9 @@ class CommandViewSet(JMSBulkModelViewSet):
             raise StorageInvalid
         else:
             qs = storage.get_command_queryset()
+        if asset_id:
+            session_ids = Session.objects.filter(asset_id=asset_id).values_list('id', flat=True)
+            qs = qs.filter(session__in=list(session_ids))
         return qs
 
     def create(self, request, *args, **kwargs):
diff --git a/apps/terminal/api/session/session.py b/apps/terminal/api/session/session.py
index fd8bc0ecd..8865801ac 100644
--- a/apps/terminal/api/session/session.py
+++ b/apps/terminal/api/session/session.py
@@ -62,7 +62,7 @@ class SessionFilterSet(BaseFilterSet):
     class Meta:
         model = Session
         fields = [
-            "user", "user_id", "asset", "account", "remote_addr",
+            "user", "user_id", "asset", "asset_id", "account", "remote_addr",
             "protocol", "is_finished", 'login_from', 'terminal'
         ]
 
diff --git a/apps/terminal/filters.py b/apps/terminal/filters.py
index ae11733cf..40c048b28 100644
--- a/apps/terminal/filters.py
+++ b/apps/terminal/filters.py
@@ -13,11 +13,12 @@ class CommandFilter(filters.FilterSet):
     user = filters.CharFilter(lookup_expr='startswith')
     input = filters.CharFilter(lookup_expr='icontains')
     asset = filters.CharFilter(field_name='asset', lookup_expr='icontains')
+    asset_id = filters.UUIDFilter(method='do_nothing')
 
     class Meta:
         model = Command
         fields = [
-            'asset', 'account', 'user', 'session', 'risk_level', 'input',
+            'asset', 'asset_id', 'account', 'user', 'session', 'risk_level', 'input',
             'date_from', 'date_to', 'session_id', 'risk_level', 'command_storage_id',
         ]