From 35a10fdd62a5b7f94deb8a0995c5c9bfa470d4a2 Mon Sep 17 00:00:00 2001
From: xinwen <coderWen@126.com>
Date: Fri, 24 Dec 2021 17:23:09 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=91=BD=E4=BB=A4=E8=AE=B0=E5=BD=95?=
 =?UTF-8?q?=E6=8E=A5=E5=8F=A3=E5=A2=9E=E5=8A=A0=20remote=5Faddr=20?=
 =?UTF-8?q?=E5=AD=97=E6=AE=B5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/terminal/api/command.py                  | 13 ++++++++++++-
 apps/terminal/backends/command/models.py      | 10 ++++++++++
 apps/terminal/backends/command/serializers.py |  1 +
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/apps/terminal/api/command.py b/apps/terminal/api/command.py
index 8b66652f1..04b3450fd 100644
--- a/apps/terminal/api/command.py
+++ b/apps/terminal/api/command.py
@@ -9,7 +9,7 @@ from rest_framework.fields import DateTimeField
 from rest_framework.response import Response
 from django.template import loader
 
-from terminal.models import CommandStorage
+from terminal.models import CommandStorage, Session
 from terminal.filters import CommandFilter
 from orgs.utils import current_org
 from common.permissions import IsOrgAdminOrAppUser, IsOrgAuditor, IsAppUser
@@ -146,15 +146,26 @@ class CommandViewSet(JMSBulkModelViewSet):
 
         page = self.paginate_queryset(queryset)
         if page is not None:
+            page = self.load_remote_addr(page)
             serializer = self.get_serializer(page, many=True)
             return self.get_paginated_response(serializer.data)
 
         # 适配像 ES 这种没有指定分页只返回少量数据的情况
         queryset = queryset[:]
 
+        queryset = self.load_remote_addr(queryset)
         serializer = self.get_serializer(queryset, many=True)
         return Response(serializer.data)
 
+    def load_remote_addr(self, queryset):
+        commands = list(queryset)
+        session_ids = {command.session for command in commands}
+        sessions = Session.objects.filter(id__in=session_ids).values_list('id', 'remote_addr')
+        session_addr_map = {str(i): addr for i, addr in sessions}
+        for command in commands:
+            command.remote_addr = session_addr_map.get(command.session, '')
+        return commands
+
     def get_queryset(self):
         command_storage_id = self.request.query_params.get('command_storage_id')
         storage = CommandStorage.objects.get(id=command_storage_id)
diff --git a/apps/terminal/backends/command/models.py b/apps/terminal/backends/command/models.py
index 7499322a7..02a16d353 100644
--- a/apps/terminal/backends/command/models.py
+++ b/apps/terminal/backends/command/models.py
@@ -4,6 +4,7 @@ import uuid
 from django.db import models
 from django.utils.translation import ugettext_lazy as _
 
+from common.utils.common import lazyproperty
 from orgs.mixins.models import OrgModelMixin
 
 
@@ -27,6 +28,15 @@ class AbstractSessionCommand(OrgModelMixin):
     class Meta:
         abstract = True
 
+    @lazyproperty
+    def remote_addr(self):
+        from terminal.models import Session
+        session = Session.objects.filter(id=self.session).first()
+        if session:
+            return session.remote_addr
+        else:
+            return ''
+
     @classmethod
     def get_risk_level_str(cls, risk_level):
         risk_mapper = dict(cls.RISK_LEVEL_CHOICES)
diff --git a/apps/terminal/backends/command/serializers.py b/apps/terminal/backends/command/serializers.py
index f435e2371..8114ddbb1 100644
--- a/apps/terminal/backends/command/serializers.py
+++ b/apps/terminal/backends/command/serializers.py
@@ -19,6 +19,7 @@ class SessionCommandSerializer(serializers.Serializer):
     risk_level_display = serializers.SerializerMethodField(label=_('Risk level display'))
     org_id = serializers.CharField(max_length=36, required=False, default='', allow_null=True, allow_blank=True)
     timestamp = serializers.IntegerField(label=_('Timestamp'))
+    remote_addr = serializers.CharField(read_only=True, label=_('Remote Address'))
 
     @staticmethod
     def get_risk_level_display(obj):