import json import logging from typing import Union from django.http import HttpResponse from rest_framework.decorators import action from rest_framework.mixins import RetrieveModelMixin, ListModelMixin, CreateModelMixin, UpdateModelMixin, \ DestroyModelMixin from rest_framework.response import Response from rest_framework.viewsets import GenericViewSet from core.utils import httpFileResponse from devices.enums import DeviceType from devices.models.endpoint_device import EndpointModel from devices.serializers.endpoint_serializers import EndpointDeviceSerializersAll from devices.services.endpoint.endpoint_antivirus import EndpointAntivirusService from devices.services.endpoint.endpoint_get_status import get_status from devices.services.endpoint.endpoint_services import EndpointManagementService, EndpointKepAliveService, \ EndpointDownloadConfigService, EndpointUploadConfigService, EndpointUpdateService _log = logging.getLogger(__name__) class EndpointDeviceApiViewSet(ListModelMixin, CreateModelMixin, RetrieveModelMixin, UpdateModelMixin, DestroyModelMixin, GenericViewSet): """ViewSet for working with endpoint device.""" queryset = EndpointModel.objects.all() serializer_class = EndpointDeviceSerializersAll def create(self, request, *args, **kwargs) -> Response: serializer = self.get_serializer(data=request.data) if not serializer.is_valid(): return Response(serializer.errors, 400) instance = serializer.save(type=DeviceType.ENDPOINT) EndpointManagementService(instance).create() return Response(self.get_serializer(instance).data, 201) def update(self, request, *args, **kwargs) -> Response: partial = kwargs.pop('partial', False) instance = self.get_object() serializer = self.get_serializer(instance, data=request.data, partial=partial) if not serializer.is_valid(): return Response(serializer.errors, 400) instance = serializer.save(settings_changed=True, ) instance.save() if getattr(instance, '_prefetched_objects_cache', None): instance._prefetched_objects_cache = {} EndpointManagementService(instance).update() return Response(self.get_serializer(instance).data) def destroy(self, request, *args, **kwargs) -> Response: instance = self.get_object() EndpointManagementService(instance).destroy() self.perform_destroy(instance) return Response({}, 204) @action(detail=True, methods=['GET', 'POST'], name="Keepalive") def keepalive(self, request, pk=None) -> Response: try: service = EndpointKepAliveService(int(pk), request.body) except EndpointModel.DoesNotExist: return Response({'status': 'error', 'reason': 'no such endpoint record'}) except json.JSONDecodeError: return Response({'status': 'error', 'error_message': 'json decode error'}) else: response = service.get_response() return Response(response) @action(detail=True, methods=['GET'], name='Status') def status(self, request, pk=None) -> Response: """ Function to respond with current Endpoint states. Current response states are: config_errors, request_config """ return Response(get_status(request, pk)) @action(detail=True, methods=['GET'], name='Download') def download(self, request, pk=None) -> Union[Response, HttpResponse]: """API for download endpoint config from MC to endpoint.""" render_format = request.accepted_renderer.format service = EndpointDownloadConfigService(int(pk)) if render_format == 'api': data, filename = service.download_as_file() response = HttpResponse(data, content_type='application/file', ) response['Content-Disposition'] = f'attachment; filename="{filename}"' return response data = service.download() service.setup_endpoint() return Response(data, status=200) @action(detail=True, methods=['POST'], name='Upload') def upload(self, request, pk=None): """API for upload endpoint config form endpoint to MC""" try: service = EndpointUploadConfigService(int(pk), request.body) except EndpointModel.DoesNotExist: return Response({'status': 'error', 'reason': 'no such endpoint record'}) except json.JSONDecodeError: return Response({'status': 'error', 'error_message': 'json decode error'}) else: response = service.upload() return Response(response) @action(detail=True, methods=["GET"], name="Endpoint_config_request") def config_request(self, request, pk=None) -> Response: """API to set the Endpoint flag in True for update MC Endpoint data in models from Endpoint instance""" service = EndpointUpdateService(int(pk)) data = service.update() _log.info(f'Request update config [{pk}] from Endpoint') return Response(data) @action(detail=True, methods=["GET"], name="Antivirus_update") def antivirus_update(self, request, pk=None) -> Union[Response, httpFileResponse]: """API to update endpoint antivirus database """ _log.info(f'Request antivirus update [{pk}] from Endpoint') service = EndpointAntivirusService(int(pk)) return service.update()