# ~*~ coding: utf-8 ~*~ # from orgs.utils import set_to_root_org __all__ = [ 'AssetsFilterMixin', 'ChangeOrgIfNeedMixin', ] class AssetsFilterMixin(object): """ 对资产进行过滤(查询,排序) """ def filter_queryset(self, queryset): queryset = self.search_assets(queryset) queryset = self.filter_labels(queryset) queryset = self.sort_assets(queryset) return queryset def search_assets(self, queryset): from perms.utils import is_obj_attr_has value = self.request.query_params.get('search') if not value: return queryset queryset = [asset for asset in queryset if is_obj_attr_has(asset, value)] return queryset def sort_assets(self, queryset): from perms.utils import sort_assets order_by = self.request.query_params.get('order') if not order_by: order_by = 'hostname' if order_by.startswith('-'): order_by = order_by.lstrip('-') reverse = True else: reverse = False queryset = sort_assets(queryset, order_by=order_by, reverse=reverse) return queryset def filter_labels(self, queryset): from assets.models import Label query_keys = self.request.query_params.keys() all_label_keys = Label.objects.values_list('name', flat=True) valid_keys = set(all_label_keys) & set(query_keys) labels_query = {} for key in valid_keys: labels_query[key] = self.request.query_params.get(key) if not labels_query: return queryset labels = set() for k, v in labels_query.items(): label = Label.objects.filter(name=k, value=v).first() if not label: continue labels.add(label) _queryset = [] for asset in queryset: _labels = set(asset.labels.all()) & set(labels) if _labels and len(_labels) == len(set(labels)): _queryset.append(asset) return _queryset class ChangeOrgIfNeedMixin(object): @staticmethod def change_org_if_need(request, kwargs): if request.user.is_authenticated and request.user.is_superuser \ or request.user.is_app \ or kwargs.get('pk') is None: set_to_root_org() def get(self, request, *args, **kwargs): self.change_org_if_need(request, kwargs) return super().get(request, *args, **kwargs)