Revert "Revert "perf:Stored command records in ES support accurate searching.""

This reverts commit 6bc1c5bd50.

apps/common/plugins/es.py

@@ -14,6 +14,7 @@ from uuid import UUID
 
 from django.utils.translation import gettext_lazy as _
 from django.db.models import QuerySet as DJQuerySet
+from django.db.models import Q
 from elasticsearch7 import Elasticsearch
 from elasticsearch7.helpers import bulk
 from elasticsearch7.exceptions import RequestError, SSLError
@@ -78,6 +79,10 @@ class ESClientV7(ESClientBase):
     def get_sort(cls, field, direction):
         return f'{field}:{direction}'
 
+    @classmethod
+    def get_sorts(cls, sorts: list):
+        return ','.join(sorts)
+
 
 class ESClientV6(ESClientV7):
 
@@ -99,6 +104,10 @@ class ESClientV8(ESClientBase):
     def get_sort(cls, field, direction):
         return {field: {'order': direction}}
 
+    @classmethod
+    def get_sorts(cls, sorts: list):
+        return sorts
+
 
 def get_es_client_version(**kwargs):
     try:
@@ -190,8 +199,7 @@ class ES(object):
                 mappings['aliases'] = {
                     self.query_index: {}
                 }
-            if self.es.indices.exists(index=self.index):
-                return
+
             try:
                 self.es.indices.create(index=self.index, body=mappings)
             except (RequestError, BadRequestError) as e:
@@ -245,6 +253,7 @@ class ES(object):
         }
         if sort is not None:
             search_params['sort'] = sort
+        logger.info('search_params: {}'.format(search_params))
         data = self.es.search(**search_params)
 
         source_data = []
@@ -319,10 +328,12 @@ class ES(object):
         kwargs = new_kwargs
 
         index_in_field = 'id__in'
+        keyword_fields = self.keyword_fields
         exact_fields = self.exact_fields
         match_fields = self.match_fields
 
         match = {}
+        search = []
         exact = {}
         index = {}
 
@@ -330,11 +341,19 @@ class ES(object):
             index['values'] = kwargs[index_in_field]
 
         for k, v in kwargs.items():
-            if k in exact_fields:
+            if k in keyword_fields:
                 exact[k] = v
+            elif k in exact_fields:
+                exact['{}.keyword'.format(k)] = v
             elif k in match_fields:
                 match[k] = v
 
+        args = kwargs.get('search')
+        for item in args:
+            for k, v in item.items():
+                if k in match_fields:
+                    search.append(item)
+
         # 处理时间
         time_field_name, time_range = self.handler_time_field(kwargs)
 
@@ -363,10 +382,12 @@ class ES(object):
         body = {
             'query': {
                 'bool': {
-                    'must': [
+                    'must': [],
+                    'should': should + [
                         {'match': {k: v}} for k, v in match.items()
+                    ] + [
+                        {'match': item} for item in search
                     ],
-                    'should': should,
                     'filter': self.handle_exact_fields(exact) +
                               [
                                   {
@@ -403,6 +424,17 @@ class QuerySet(DJQuerySet):
         _method_calls = {k: list(v) for k, v in groupby(self._method_calls, lambda x: x[0])}
         return _method_calls
 
+    def _grouped_search_args(self, query):
+        conditions = {}
+        for q in query:
+            for c in q.children:
+                if isinstance(c, Q):
+                    child = self._grouped_search_args(c)
+                    [conditions.setdefault(k, []).extend(v) for k, v in child.items()]
+                else:
+                    conditions.setdefault(c[0], []).append(c[1])
+        return conditions
+
     @lazyproperty
     def _filter_kwargs(self):
         _method_calls = self._grouped_method_calls
@@ -410,14 +442,14 @@ class QuerySet(DJQuerySet):
         if not filter_calls:
             return {}
         names, multi_args, multi_kwargs = zip(*filter_calls)
-        args = {
-            key: value
-            for arg in multi_args if arg
-            for key, value in arg[0].children
-        }
+
+        # input 输入
+        multi_args = tuple(reduce(lambda x, y: x + y, (sub for sub in multi_args if sub),()))
+        args = self._grouped_search_args(multi_args)
+        striped_args = [{k.replace('__icontains', ''): v} for k, values in args.items() for v in values]
+
         kwargs = reduce(lambda x, y: {**x, **y}, multi_kwargs, {})
-        kwargs.update(args)
-        striped_kwargs = {}
+        striped_kwargs = {'search': striped_args}
         for k, v in kwargs.items():
             k = k.replace('__exact', '')
             k = k.replace('__startswith', '')
@@ -428,6 +460,7 @@ class QuerySet(DJQuerySet):
     @lazyproperty
     def _sort(self):
         order_by = self._grouped_method_calls.get('order_by')
+        _sorts = [self._storage.client.get_sort('_score', 'desc')]
         if order_by:
             for call in reversed(order_by):
                 fields = call[1]
@@ -440,7 +473,10 @@ class QuerySet(DJQuerySet):
                         direction = 'asc'
                     field = field.lstrip('-+')
                     sort = self._storage.client.get_sort(field, direction)
-                    return sort
+                    _sorts.append(sort)
+                    break
+        sorts = self._storage.client.get_sorts(_sorts)
+        return sorts
 
     def __execute(self):
         _filter_kwargs = self._filter_kwargs
@@ -514,4 +550,4 @@ class QuerySet(DJQuerySet):
         return iter(self.__execute())
 
     def __len__(self):
-        return self.count()
+        return self.count()

apps/terminal/backends/command/es.py

@@ -27,8 +27,8 @@ class CommandStore(ES):
                 "type": "long"
             }
         }
-        exact_fields = {}
-        match_fields = {'input', 'risk_level', 'user', 'asset', 'system_user'}
+        exact_fields = {'input', 'risk_level', 'user', 'asset', 'account'}
+        match_fields = {'input'}
         keyword_fields = {'session', 'org_id'}
 
         super().__init__(config, properties, keyword_fields, exact_fields, match_fields)