perf: update status

pull/14774/head^2
ibuler 2024-12-11 09:59:59 +08:00
parent c95c3099b7
commit 8aeffba64d
2 changed files with 191 additions and 120 deletions

View File

@ -11,31 +11,39 @@ from .models import Account, GatheredAccount, ChangeSecretRecord
class AccountFilterSet(BaseFilterSet): class AccountFilterSet(BaseFilterSet):
ip = drf_filters.CharFilter(field_name='address', lookup_expr='exact') ip = drf_filters.CharFilter(field_name="address", lookup_expr="exact")
hostname = drf_filters.CharFilter(field_name='name', lookup_expr='exact') hostname = drf_filters.CharFilter(field_name="name", lookup_expr="exact")
username = drf_filters.CharFilter(field_name="username", lookup_expr='exact') username = drf_filters.CharFilter(field_name="username", lookup_expr="exact")
address = drf_filters.CharFilter(field_name="asset__address", lookup_expr='exact') address = drf_filters.CharFilter(field_name="asset__address", lookup_expr="exact")
asset_id = drf_filters.CharFilter(field_name="asset", lookup_expr='exact') asset_id = drf_filters.CharFilter(field_name="asset", lookup_expr="exact")
asset = drf_filters.CharFilter(field_name='asset', lookup_expr='exact') asset = drf_filters.CharFilter(field_name="asset", lookup_expr="exact")
assets = drf_filters.CharFilter(field_name='asset_id', lookup_expr='exact') assets = drf_filters.CharFilter(field_name="asset_id", lookup_expr="exact")
nodes = drf_filters.CharFilter(method='filter_nodes') nodes = drf_filters.CharFilter(method="filter_nodes")
node_id = drf_filters.CharFilter(method='filter_nodes') node_id = drf_filters.CharFilter(method="filter_nodes")
has_secret = drf_filters.BooleanFilter(method='filter_has_secret') has_secret = drf_filters.BooleanFilter(method="filter_has_secret")
platform = drf_filters.CharFilter(field_name='asset__platform_id', lookup_expr='exact') platform = drf_filters.CharFilter(
category = drf_filters.CharFilter(field_name='asset__platform__category', lookup_expr='exact') field_name="asset__platform_id", lookup_expr="exact"
type = drf_filters.CharFilter(field_name='asset__platform__type', lookup_expr='exact') )
latest_discovery = drf_filters.BooleanFilter(method='filter_latest') category = drf_filters.CharFilter(
latest_accessed = drf_filters.BooleanFilter(method='filter_latest') field_name="asset__platform__category", lookup_expr="exact"
latest_updated = drf_filters.BooleanFilter(method='filter_latest') )
latest_secret_changed = drf_filters.BooleanFilter(method='filter_latest') type = drf_filters.CharFilter(
latest_secret_change_failed = drf_filters.BooleanFilter(method='filter_latest') field_name="asset__platform__type", lookup_expr="exact"
risk = drf_filters.CharFilter(method='filter_risk', ) )
long_time_no_change_secret = drf_filters.BooleanFilter(method='filter_long_time') latest_discovery = drf_filters.BooleanFilter(method="filter_latest")
long_time_no_verified = drf_filters.BooleanFilter(method='filter_long_time') latest_accessed = drf_filters.BooleanFilter(method="filter_latest")
latest_updated = drf_filters.BooleanFilter(method="filter_latest")
latest_secret_changed = drf_filters.BooleanFilter(method="filter_latest")
latest_secret_change_failed = drf_filters.BooleanFilter(method="filter_latest")
risk = drf_filters.CharFilter(
method="filter_risk",
)
long_time_no_change_secret = drf_filters.BooleanFilter(method="filter_long_time")
long_time_no_verified = drf_filters.BooleanFilter(method="filter_long_time")
@staticmethod @staticmethod
def filter_has_secret(queryset, name, has_secret): def filter_has_secret(queryset, name, has_secret):
q = Q(_secret__isnull=True) | Q(_secret='') q = Q(_secret__isnull=True) | Q(_secret="")
if has_secret: if has_secret:
return queryset.exclude(q) return queryset.exclude(q)
else: else:
@ -45,15 +53,15 @@ class AccountFilterSet(BaseFilterSet):
def filter_long_time(queryset, name, value): def filter_long_time(queryset, name, value):
date = timezone.now() - timezone.timedelta(days=30) date = timezone.now() - timezone.timedelta(days=30)
if name == 'long_time_no_change_secret': if name == "long_time_no_change_secret":
field = 'date_change_secret' field = "date_change_secret"
confirm_field = 'change_secret_status' confirm_field = "change_secret_status"
else: else:
field = 'date_verified' field = "date_verified"
confirm_field = 'connectivity' confirm_field = "connectivity"
q = Q(**{f'{field}__lt': date}) | Q(**{f'{field}__isnull': True}) q = Q(**{f"{field}__lt": date}) | Q(**{f"{field}__isnull": True})
confirm_q = {f'{confirm_field}': 'na'} confirm_q = {f"{confirm_field}": "na"}
queryset = queryset.exclude(**confirm_q).filter(q) queryset = queryset.exclude(**confirm_q).filter(q)
return queryset return queryset
@ -62,9 +70,11 @@ class AccountFilterSet(BaseFilterSet):
if not value: if not value:
return queryset return queryset
queryset = queryset.prefetch_related('risks') \ queryset = (
.annotate(risk=F('risks__risk'), confirmed=F('risks__confirmed')) \ queryset.prefetch_related("risks")
.annotate(risk=F("risks__risk"), confirmed=F("risks__confirmed"))
.filter(risk=value, confirmed=False) .filter(risk=value, confirmed=False)
)
return queryset return queryset
@staticmethod @staticmethod
@ -75,17 +85,19 @@ class AccountFilterSet(BaseFilterSet):
date = timezone.now() - timezone.timedelta(days=7) date = timezone.now() - timezone.timedelta(days=7)
kwargs = {} kwargs = {}
if name == 'latest_discovery': if name == "latest_discovery":
kwargs.update({'date_created__gte': date, 'source': 'collected'}) kwargs.update({"date_created__gte": date, "source": "collected"})
elif name == 'latest_accessed': elif name == "latest_accessed":
kwargs.update({'date_last_login__gte': date}) kwargs.update({"date_last_login__gte": date})
elif name == 'latest_updated': elif name == "latest_updated":
kwargs.update({'date_updated__gte': date}) kwargs.update({"date_updated__gte": date})
elif name == 'latest_secret_changed': elif name == "latest_secret_changed":
kwargs.update({'date_change_secret__gt': date}) kwargs.update({"date_change_secret__gt": date})
if name == 'latest_secret_change_failed': if name == "latest_secret_change_failed":
queryset = queryset.filter(date_change_secret__gt=date).exclude(change_secret_status='ok') queryset = queryset.filter(date_change_secret__gt=date).exclude(
change_secret_status="ok"
)
if kwargs: if kwargs:
queryset = queryset.filter(date_last_login__gte=date) queryset = queryset.filter(date_last_login__gte=date)
@ -100,19 +112,22 @@ class AccountFilterSet(BaseFilterSet):
node_qs = Node.objects.none() node_qs = Node.objects.none()
for node in nodes: for node in nodes:
node_qs |= node.get_all_children(with_self=True) node_qs |= node.get_all_children(with_self=True)
node_ids = list(node_qs.values_list('id', flat=True)) node_ids = list(node_qs.values_list("id", flat=True))
queryset = queryset.filter(asset__nodes__in=node_ids) queryset = queryset.filter(asset__nodes__in=node_ids)
return queryset return queryset
class Meta: class Meta:
model = Account model = Account
fields = ['id', 'asset', 'source_id', 'secret_type', 'category', 'type'] fields = ["id", "asset", "source_id", "secret_type", "category", "type"]
class GatheredAccountFilterSet(BaseFilterSet): class GatheredAccountFilterSet(BaseFilterSet):
node_id = drf_filters.CharFilter(method='filter_nodes') node_id = drf_filters.CharFilter(method="filter_nodes")
asset_id = drf_filters.CharFilter(field_name='asset_id', lookup_expr='exact') asset_id = drf_filters.CharFilter(field_name="asset_id", lookup_expr="exact")
asset_name = drf_filters.CharFilter(field_name='asset__name', lookup_expr='icontains') asset_name = drf_filters.CharFilter(
field_name="asset__name", lookup_expr="icontains"
)
status = drf_filters.CharFilter(field_name="status", lookup_expr="exact")
@staticmethod @staticmethod
def filter_nodes(queryset, name, value): def filter_nodes(queryset, name, value):
@ -120,15 +135,20 @@ class GatheredAccountFilterSet(BaseFilterSet):
class Meta: class Meta:
model = GatheredAccount model = GatheredAccount
fields = ['id', 'username'] fields = ["id", "username"]
class ChangeSecretRecordFilterSet(BaseFilterSet): class ChangeSecretRecordFilterSet(BaseFilterSet):
asset_name = drf_filters.CharFilter(field_name='asset__name', lookup_expr='icontains') asset_name = drf_filters.CharFilter(
account_username = drf_filters.CharFilter(field_name='account__username', lookup_expr='icontains') field_name="asset__name", lookup_expr="icontains"
execution_id = drf_filters.CharFilter(field_name='execution_id', lookup_expr='exact') )
account_username = drf_filters.CharFilter(
days = drf_filters.NumberFilter(method='filter_days') field_name="account__username", lookup_expr="icontains"
)
execution_id = drf_filters.CharFilter(
field_name="execution_id", lookup_expr="exact"
)
days = drf_filters.NumberFilter(method="filter_days")
@staticmethod @staticmethod
def filter_days(queryset, name, value): def filter_days(queryset, name, value):
@ -141,4 +161,4 @@ class ChangeSecretRecordFilterSet(BaseFilterSet):
class Meta: class Meta:
model = ChangeSecretRecord model = ChangeSecretRecord
fields = ['id', 'status', 'asset_id', 'execution'] fields = ["id", "status", "asset_id", "execution"]

View File

@ -4,6 +4,7 @@ import base64
import json import json
import logging import logging
from collections import defaultdict from collections import defaultdict
from django.utils import timezone
from django.core.cache import cache from django.core.cache import cache
from django.core.exceptions import ImproperlyConfigured from django.core.exceptions import ImproperlyConfigured
@ -18,18 +19,26 @@ from rest_framework.filters import OrderingFilter
from common import const from common import const
from common.db.fields import RelatedManager from common.db.fields import RelatedManager
logger = logging.getLogger('jumpserver.common') logger = logging.getLogger("jumpserver.common")
__all__ = [ __all__ = [
"DatetimeRangeFilterBackend", "IDSpmFilterBackend", "DatetimeRangeFilterBackend",
'IDInFilterBackend', "CustomFilterBackend", "IDSpmFilterBackend",
"BaseFilterSet", 'IDNotFilterBackend', "IDInFilterBackend",
'NotOrRelFilterBackend', 'LabelFilterBackend', "CustomFilterBackend",
'RewriteOrderingFilter', 'AttrRulesFilterBackend' "BaseFilterSet",
"IDNotFilterBackend",
"NotOrRelFilterBackend",
"LabelFilterBackend",
"RewriteOrderingFilter",
"AttrRulesFilterBackend",
] ]
class BaseFilterSet(drf_filters.FilterSet): class BaseFilterSet(drf_filters.FilterSet):
days = drf_filters.NumberFilter(method="filter_days")
days__lt = drf_filters.NumberFilter(method="filter_days")
def do_nothing(self, queryset, name, value): def do_nothing(self, queryset, name, value):
return queryset return queryset
@ -38,6 +47,22 @@ class BaseFilterSet(drf_filters.FilterSet):
return self.form.cleaned_data[k] return self.form.cleaned_data[k]
return default return default
@staticmethod
def filter_days(queryset, name, value):
try:
value = int(value)
except ValueError:
return queryset.none()
if name == 'days':
arg = 'date_created__gte'
else:
arg = 'date_created__lt'
date = timezone.now() - timezone.timedelta(days=value)
kwargs = {arg: date}
return queryset.filter(**kwargs)
class DatetimeRangeFilterBackend(filters.BaseFilterBackend): class DatetimeRangeFilterBackend(filters.BaseFilterBackend):
def get_schema_fields(self, view): def get_schema_fields(self, view):
@ -50,18 +75,20 @@ class DatetimeRangeFilterBackend(filters.BaseFilterBackend):
for v in date_range_keyword: for v in date_range_keyword:
ret.append( ret.append(
coreapi.Field( coreapi.Field(
name=v, location='query', required=False, type='string', name=v,
location="query",
required=False,
type="string",
schema=coreschema.String( schema=coreschema.String(
title=v, title=v, description="%s %s" % (attr, v)
description='%s %s' % (attr, v) ),
)
) )
) )
return ret return ret
def _get_date_range_filter_fields(self, view): def _get_date_range_filter_fields(self, view):
if not hasattr(view, 'date_range_filter_fields'): if not hasattr(view, "date_range_filter_fields"):
return {} return {}
try: try:
return dict(view.date_range_filter_fields) return dict(view.date_range_filter_fields)
@ -75,7 +102,9 @@ class DatetimeRangeFilterBackend(filters.BaseFilterBackend):
('db column', ('query param date from', 'query param date to')) ('db column', ('query param date from', 'query param date to'))
] ]
``` ```
""".format(view.name) """.format(
view.name
)
logger.error(msg) logger.error(msg)
raise ImproperlyConfigured(msg) raise ImproperlyConfigured(msg)
@ -110,14 +139,17 @@ class IDSpmFilterBackend(filters.BaseFilterBackend):
def get_schema_fields(self, view): def get_schema_fields(self, view):
return [ return [
coreapi.Field( coreapi.Field(
name='spm', location='query', required=False, name="spm",
type='string', example='', location="query",
description='Pre post objects id get spm id, then using filter' required=False,
type="string",
example="",
description="Pre post objects id get spm id, then using filter",
) )
] ]
def filter_queryset(self, request, queryset, view): def filter_queryset(self, request, queryset, view):
spm = request.query_params.get('spm') spm = request.query_params.get("spm")
if not spm: if not spm:
return queryset return queryset
cache_key = const.KEY_CACHE_RESOURCE_IDS.format(spm) cache_key = const.KEY_CACHE_RESOURCE_IDS.format(spm)
@ -127,7 +159,7 @@ class IDSpmFilterBackend(filters.BaseFilterBackend):
return queryset.none() return queryset.none()
if isinstance(resource_ids, str): if isinstance(resource_ids, str):
resource_ids = [resource_ids] resource_ids = [resource_ids]
if hasattr(view, 'filter_spm_queryset'): if hasattr(view, "filter_spm_queryset"):
queryset = view.filter_spm_queryset(resource_ids, queryset) queryset = view.filter_spm_queryset(resource_ids, queryset)
else: else:
queryset = queryset.filter(id__in=resource_ids) queryset = queryset.filter(id__in=resource_ids)
@ -138,17 +170,20 @@ class IDInFilterBackend(filters.BaseFilterBackend):
def get_schema_fields(self, view): def get_schema_fields(self, view):
return [ return [
coreapi.Field( coreapi.Field(
name='ids', location='query', required=False, name="ids",
type='string', example='/api/v1/users/users?ids=1,2,3', location="query",
description='Filter by id set' required=False,
type="string",
example="/api/v1/users/users?ids=1,2,3",
description="Filter by id set",
) )
] ]
def filter_queryset(self, request, queryset, view): def filter_queryset(self, request, queryset, view):
ids = request.query_params.get('ids') ids = request.query_params.get("ids")
if not ids: if not ids:
return queryset return queryset
id_list = [i.strip() for i in ids.split(',')] id_list = [i.strip() for i in ids.split(",")]
queryset = queryset.filter(id__in=id_list) queryset = queryset.filter(id__in=id_list)
return queryset return queryset
@ -157,17 +192,20 @@ class IDNotFilterBackend(filters.BaseFilterBackend):
def get_schema_fields(self, view): def get_schema_fields(self, view):
return [ return [
coreapi.Field( coreapi.Field(
name='id!', location='query', required=False, name="id!",
type='string', example='/api/v1/users/users?id!=1,2,3', location="query",
description='Exclude by id set' required=False,
type="string",
example="/api/v1/users/users?id!=1,2,3",
description="Exclude by id set",
) )
] ]
def filter_queryset(self, request, queryset, view): def filter_queryset(self, request, queryset, view):
ids = request.query_params.get('id!') ids = request.query_params.get("id!")
if not ids: if not ids:
return queryset return queryset
id_list = [i.strip() for i in ids.split(',')] id_list = [i.strip() for i in ids.split(",")]
queryset = queryset.exclude(id__in=id_list) queryset = queryset.exclude(id__in=id_list)
return queryset return queryset
@ -176,26 +214,30 @@ class LabelFilterBackend(filters.BaseFilterBackend):
def get_schema_fields(self, view): def get_schema_fields(self, view):
return [ return [
coreapi.Field( coreapi.Field(
name='label', location='query', required=False, name="label",
type='string', example='/api/v1/users/users?label=abc', location="query",
description='Filter by label' required=False,
type="string",
example="/api/v1/users/users?label=abc",
description="Filter by label",
) )
] ]
@staticmethod @staticmethod
def parse_labels(labels_id): def parse_labels(labels_id):
from labels.models import Label from labels.models import Label
label_ids = [i.strip() for i in labels_id.split(',')]
label_ids = [i.strip() for i in labels_id.split(",")]
cleaned = [] cleaned = []
args = [] args = []
for label_id in label_ids: for label_id in label_ids:
kwargs = {} kwargs = {}
if ':' in label_id: if ":" in label_id:
k, v = label_id.split(':', 1) k, v = label_id.split(":", 1)
kwargs['name'] = k.strip() kwargs["name"] = k.strip()
if v != '*': if v != "*":
kwargs['value'] = v.strip() kwargs["value"] = v.strip()
args.append(kwargs) args.append(kwargs)
else: else:
cleaned.append(label_id) cleaned.append(label_id)
@ -209,14 +251,14 @@ class LabelFilterBackend(filters.BaseFilterBackend):
return cleaned return cleaned
def filter_queryset(self, request, queryset, view): def filter_queryset(self, request, queryset, view):
labels_id = request.query_params.get('labels') labels_id = request.query_params.get("labels")
if not labels_id: if not labels_id:
return queryset return queryset
if not hasattr(queryset, 'model'): if not hasattr(queryset, "model"):
return queryset return queryset
if not hasattr(queryset.model, 'label_model'): if not hasattr(queryset.model, "label_model"):
return queryset return queryset
model = queryset.model.label_model() model = queryset.model.label_model()
@ -225,7 +267,8 @@ class LabelFilterBackend(filters.BaseFilterBackend):
model_name = model._meta.model_name model_name = model._meta.model_name
full_resources = labeled_resource_cls.objects.filter( full_resources = labeled_resource_cls.objects.filter(
res_type__app_label=app_label, res_type__model=model_name, res_type__app_label=app_label,
res_type__model=model_name,
) )
labels = self.parse_labels(labels_id) labels = self.parse_labels(labels_id)
grouped = defaultdict(set) grouped = defaultdict(set)
@ -234,8 +277,10 @@ class LabelFilterBackend(filters.BaseFilterBackend):
matched_ids = set() matched_ids = set()
for name, label_ids in grouped.items(): for name, label_ids in grouped.items():
resources = model.filter_resources_by_labels(full_resources, label_ids, rel='any') resources = model.filter_resources_by_labels(
res_ids = resources.values_list('res_id', flat=True) full_resources, label_ids, rel="any"
)
res_ids = resources.values_list("res_id", flat=True)
if not matched_ids: if not matched_ids:
matched_ids = set(res_ids) matched_ids = set(res_ids)
else: else:
@ -249,16 +294,14 @@ class CustomFilterBackend(filters.BaseFilterBackend):
def get_schema_fields(self, view): def get_schema_fields(self, view):
fields = [] fields = []
defaults = dict( defaults = dict(
location='query', required=False, location="query", required=False, type="string", example="", description=""
type='string', example='',
description=''
) )
if not hasattr(view, 'custom_filter_fields'): if not hasattr(view, "custom_filter_fields"):
return [] return []
for field in view.custom_filter_fields: for field in view.custom_filter_fields:
if isinstance(field, str): if isinstance(field, str):
defaults['name'] = field defaults["name"] = field
elif isinstance(field, dict): elif isinstance(field, dict):
defaults.update(field) defaults.update(field)
else: else:
@ -270,7 +313,7 @@ class CustomFilterBackend(filters.BaseFilterBackend):
return queryset return queryset
def current_user_filter(user_field='user'): def current_user_filter(user_field="user"):
class CurrentUserFilter(filters.BaseFilterBackend): class CurrentUserFilter(filters.BaseFilterBackend):
def filter_queryset(self, request, queryset, view): def filter_queryset(self, request, queryset, view):
return queryset.filter(**{user_field: request.user}) return queryset.filter(**{user_field: request.user})
@ -290,27 +333,30 @@ class AttrRulesFilterBackend(filters.BaseFilterBackend):
def get_schema_fields(self, view): def get_schema_fields(self, view):
return [ return [
coreapi.Field( coreapi.Field(
name='attr_rules', location='query', required=False, name="attr_rules",
type='string', example='/api/v1/users/users?attr_rules=jsonbase64', location="query",
description='Filter by json like {"type": "attrs", "attrs": []} to base64' required=False,
type="string",
example="/api/v1/users/users?attr_rules=jsonbase64",
description='Filter by json like {"type": "attrs", "attrs": []} to base64',
) )
] ]
def filter_queryset(self, request, queryset, view): def filter_queryset(self, request, queryset, view):
attr_rules = request.query_params.get('attr_rules') attr_rules = request.query_params.get("attr_rules")
if not attr_rules: if not attr_rules:
return queryset return queryset
try: try:
attr_rules = base64.b64decode(attr_rules.encode('utf-8')) attr_rules = base64.b64decode(attr_rules.encode("utf-8"))
except Exception: except Exception:
raise ValidationError({'attr_rules': 'attr_rules should be base64'}) raise ValidationError({"attr_rules": "attr_rules should be base64"})
try: try:
attr_rules = json.loads(attr_rules) attr_rules = json.loads(attr_rules)
except Exception: except Exception:
raise ValidationError({'attr_rules': 'attr_rules should be json'}) raise ValidationError({"attr_rules": "attr_rules should be json"})
logger.debug('attr_rules: %s', attr_rules) logger.debug("attr_rules: %s", attr_rules)
qs = RelatedManager.get_to_filter_qs(attr_rules, queryset.model) qs = RelatedManager.get_to_filter_qs(attr_rules, queryset.model)
for q in qs: for q in qs:
queryset = queryset.filter(q) queryset = queryset.filter(q)
@ -321,33 +367,38 @@ class NotOrRelFilterBackend(filters.BaseFilterBackend):
def get_schema_fields(self, view): def get_schema_fields(self, view):
return [ return [
coreapi.Field( coreapi.Field(
name='_rel', location='query', required=False, name="_rel",
type='string', example='/api/v1/users/users?name=abc&username=def&_rel=union', location="query",
description='Filter by rel, or not, default is and' required=False,
type="string",
example="/api/v1/users/users?name=abc&username=def&_rel=union",
description="Filter by rel, or not, default is and",
) )
] ]
def filter_queryset(self, request, queryset, view): def filter_queryset(self, request, queryset, view):
_rel = request.query_params.get('_rel') _rel = request.query_params.get("_rel")
if not _rel or _rel not in ('or', 'not'): if not _rel or _rel not in ("or", "not"):
return queryset return queryset
if _rel == 'not': if _rel == "not":
queryset.query.where.negated = True queryset.query.where.negated = True
elif _rel == 'or': elif _rel == "or":
queryset.query.where.connector = 'OR' queryset.query.where.connector = "OR"
queryset._result_cache = None queryset._result_cache = None
return queryset return queryset
class RewriteOrderingFilter(OrderingFilter): class RewriteOrderingFilter(OrderingFilter):
default_ordering_if_has = ('name', ) default_ordering_if_has = ("name",)
def get_default_ordering(self, view): def get_default_ordering(self, view):
ordering = super().get_default_ordering(view) ordering = super().get_default_ordering(view)
# 如果 view.ordering = [] 表示不排序, 这样可以节约性能 (比如: 用户授权的资产) # 如果 view.ordering = [] 表示不排序, 这样可以节约性能 (比如: 用户授权的资产)
if ordering is not None: if ordering is not None:
return ordering return ordering
ordering_fields = getattr(view, 'ordering_fields', self.ordering_fields) ordering_fields = getattr(view, "ordering_fields", self.ordering_fields)
if ordering_fields: if ordering_fields:
ordering = tuple([f for f in ordering_fields if f in self.default_ordering_if_has]) ordering = tuple(
[f for f in ordering_fields if f in self.default_ordering_if_has]
)
return ordering return ordering