# -*- coding: utf-8 -*- # import uuid from django.utils import timezone from rest_framework import generics from rest_framework.fields import DateTimeField from rest_framework.response import Response from acls.models import CommandFilterACL, CommandGroup from common.api import JMSBulkModelViewSet from common.utils import get_logger from orgs.utils import current_org from terminal.backends import ( get_command_storage, get_multi_command_storage ) from terminal.const import RiskLevelChoices from terminal.exceptions import StorageInvalid from terminal.filters import CommandFilter from terminal.models import CommandStorage, Session, Command from terminal.notifications import CommandAlertMessage, CommandWarningMessage from terminal.serializers import ( SessionCommandSerializer, InsecureCommandAlertSerializer ) logger = get_logger(__name__) __all__ = ['CommandViewSet', 'InsecureCommandAlertAPI'] class CommandQueryMixin: command_store = get_command_storage() filterset_fields = [ "asset", "account", "user", "session", "risk_level", "input" ] default_days_ago = 5 @staticmethod def get_org_id(): if current_org.is_default(): org_id = '' else: org_id = current_org.id return org_id def get_query_risk_level(self): risk_level = self.request.query_params.get('risk_level') if risk_level is None: return None if risk_level.isdigit(): return int(risk_level) return None def get_queryset(self): # 解决访问 /docs/ 问题 if hasattr(self, 'swagger_fake_view'): return self.command_store.model.objects.none() date_from, date_to = self.get_date_range() q = self.request.query_params multi_command_storage = get_multi_command_storage() queryset = multi_command_storage.filter( date_from=date_from, date_to=date_to, user=q.get("user"), asset=q.get("asset"), account=q.get("account"), input=q.get("input"), session=q.get("session_id", q.get('session')), risk_level=self.get_query_risk_level(), org_id=self.get_org_id(), ) return queryset def filter_queryset(self, queryset): # 解决es存储命令时，父类根据filter_fields过滤出现异常的问题，返回的queryset类型list return queryset def get_date_range(self): now = timezone.now() days_ago = now - timezone.timedelta(days=self.default_days_ago) date_from_st = days_ago.timestamp() date_to_st = now.timestamp() query_params = self.request.query_params date_from_q = query_params.get("date_from") date_to_q = query_params.get("date_to") dt_parser = DateTimeField().to_internal_value if date_from_q: date_from_st = dt_parser(date_from_q).timestamp() if date_to_q: date_to_st = dt_parser(date_to_q).timestamp() return date_from_st, date_to_st class CommandViewSet(JMSBulkModelViewSet): """接受app发送来的command log, 格式如下 { "user": "admin", "asset": "localhost", "account": "web", "session": "xxxxxx", "input": "whoami", "output": "d2hvbWFp", # base64.b64encode(s) "timestamp": 1485238673.0 } """ command_store = get_command_storage() serializer_class = SessionCommandSerializer filterset_class = CommandFilter model = Command search_fields = ('input',) ordering_fields = ('timestamp', 'risk_level') def merge_all_storage_list(self, request, *args, **kwargs): merged_commands = [] storages = CommandStorage.objects.all() for storage in storages: if not storage.is_valid(): continue qs = storage.get_command_queryset() commands = self.filter_queryset(qs) merged_commands.extend(commands[:]) # ES 默认只取 10 条数据 order = self.request.query_params.get('order', None) if order == 'timestamp': merged_commands.sort(key=lambda command: command.timestamp) else: merged_commands.sort(key=lambda command: command.timestamp, reverse=True) page = self.paginate_queryset(merged_commands) if page is not None: serializer = self.get_serializer(page, many=True) return self.get_paginated_response(serializer.data) serializer = self.get_serializer(merged_commands, many=True) return Response(serializer.data) def list(self, request, *args, **kwargs): command_storage_id = self.request.query_params.get('command_storage_id') session_id = self.request.query_params.get('session_id') if session_id and not command_storage_id: # 会话里的命令列表肯定会提供 session_id，这里防止 merge 的时候取全量的数据 return self.merge_all_storage_list(request, *args, **kwargs) queryset = self.get_queryset() queryset = self.filter_queryset(queryset) 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') asset_id = self.request.query_params.get('asset_id') if not command_storage_id: return Command.objects.none() storage = CommandStorage.objects.get(id=command_storage_id) if not storage.is_valid(): raise StorageInvalid else: qs = storage.get_command_queryset() if asset_id: session_ids = Session.objects.filter(asset_id=asset_id).values_list('id', flat=True) qs = qs.filter(session__in=list(session_ids)) return qs def create(self, request, *args, **kwargs): serializer = self.serializer_class(data=request.data, many=True) if serializer.is_valid(): ok = self.command_store.bulk_save(serializer.validated_data) if ok: return Response("ok", status=201) else: return Response("Save error", status=500) else: msg = "Command not valid: {}".format(serializer.errors) logger.error(msg) return Response({"msg": msg}, status=401) class InsecureCommandAlertAPI(generics.CreateAPIView): serializer_class = InsecureCommandAlertSerializer rbac_perms = { 'POST': 'terminal.add_command' } def post(self, request, *args, **kwargs): serializer = InsecureCommandAlertSerializer(data=request.data, many=True) serializer.is_valid(raise_exception=True) commands = serializer.validated_data session_ids, acl_ids, cmd_group_ids = set(), set(), set() for command in commands: session_ids.add(command.get('session')) acl_ids.add(command.get('cmd_filter_acl')) cmd_group_ids.add(command.get('cmd_group')) sessions = Session.objects.filter(id__in=session_ids).only( 'id', 'org_id', 'asset', 'asset_id', 'user', 'user_id', 'account', 'account_id' ) session_mapper = {str(i.id): i for i in sessions} acls = CommandFilterACL.objects.filter(id__in=acl_ids).only('id', 'name', 'reviewers') acl_mapper = {str(i.id): i for i in acls} cmd_groups = CommandGroup.objects.filter(id__in=cmd_group_ids).only('id', 'name') cmd_group_mapper = {str(i.id): i for i in cmd_groups} for command in commands: cmd_acl = acl_mapper.get(command['cmd_filter_acl']) command['_cmd_filter_acl'] = cmd_acl cmd_group = cmd_group_mapper.get(command['cmd_group']) command['_cmd_group'] = cmd_group session = session_mapper.get(command['session']) risk_level = command.get('risk_level') if session: command.update({ '_account': session.account, '_org_name': session.org.name, }) if risk_level in [RiskLevelChoices.reject, RiskLevelChoices.review_reject]: CommandAlertMessage(command).publish_async() elif risk_level in [RiskLevelChoices.warning]: for reviewer in cmd_acl.reviewers.all(): CommandWarningMessage(reviewer, command).publish_async() else: logger.info(f'Risk level ignore: {RiskLevelChoices.get_label(risk_level)}({risk_level})') return Response({'msg': 'ok'})