from django_filters import rest_framework as filters from django.db.models import QuerySet from orgs.utils import current_org from terminal.models import Command, CommandStorage class CommandFilter(filters.FilterSet): date_from = filters.DateTimeFilter(method='do_nothing') date_to = filters.DateTimeFilter(method='do_nothing') session_id = filters.CharFilter(field_name='session') command_storage_id = filters.UUIDFilter(method='do_nothing') user = filters.CharFilter(lookup_expr='startswith') input = filters.CharFilter(lookup_expr='icontains') class Meta: model = Command fields = [ 'asset', 'system_user', 'user', 'session', 'risk_level', 'input', 'date_from', 'date_to', 'session_id', 'risk_level', 'command_storage_id', ] def do_nothing(self, queryset, name, value): return queryset @property def qs(self): qs = super().qs qs = qs.filter(org_id=self.get_org_id()) qs = self.filter_by_timestamp(qs) return qs def filter_by_timestamp(self, qs: QuerySet): date_from = self.form.cleaned_data.get('date_from') date_to = self.form.cleaned_data.get('date_to') filters = {} if date_from: date_from = date_from.timestamp() filters['timestamp__gte'] = date_from if date_to: date_to = date_to.timestamp() filters['timestamp__lte'] = date_to qs = qs.filter(**filters) return qs @staticmethod def get_org_id(): org_id = current_org.id return org_id class CommandFilterForStorageTree(CommandFilter): asset = filters.CharFilter(method='do_nothing') system_user = filters.CharFilter(method='do_nothing') session = filters.CharFilter(method='do_nothing') risk_level = filters.NumberFilter(method='do_nothing') class Meta: model = CommandStorage fields = [ 'asset', 'system_user', 'user', 'session', 'risk_level', 'input', 'date_from', 'date_to', 'session_id', 'risk_level', 'command_storage_id', ] class CommandStorageFilter(filters.FilterSet): real = filters.BooleanFilter(method='filter_real') class Meta: model = CommandStorage fields = ['real', 'name', 'type'] def filter_real(self, queryset, name, value): if value: queryset = queryset.exclude(name='null') return queryset