119 lines
5.1 KiB
Python
119 lines
5.1 KiB
Python
import logging
|
|
|
|
from django.http import Http404
|
|
from rest_framework import status
|
|
from rest_framework.decorators import action
|
|
from rest_framework.exceptions import APIException
|
|
from rest_framework.response import Response
|
|
from rest_framework.viewsets import ModelViewSet
|
|
|
|
from devices.enums import DeviceType
|
|
from devices.models.sensor import ArmaSensor
|
|
from devices.serializers.sensor_serializers import SensorSerializer, ZeekSettingsUpdateSerializer, \
|
|
ZeekProtocolsDisableSerializer, CreateSensorSerializer
|
|
from devices.services.sensor.enums import VectorMessage, SystemMessage, ZeekMessage, SuricataMessage
|
|
from devices.services.sensor.rabbitmq import SensorManagement
|
|
from devices.services.sensor.service import SensorService
|
|
from devices.services.sensor.utils import RabbitMQUserManagement
|
|
from devices.services.vector import VectorService
|
|
|
|
_log = logging.getLogger(__name__)
|
|
|
|
|
|
class ArmaSensorViewSet(ModelViewSet):
|
|
queryset = ArmaSensor.objects.all()
|
|
|
|
def get_serializer_class(self):
|
|
if self.action == 'create':
|
|
return CreateSensorSerializer
|
|
return SensorSerializer
|
|
|
|
def perform_update(self, serializer):
|
|
|
|
sensor = serializer.save()
|
|
body = serializer.data
|
|
SensorManagement().send_message(sensor=sensor, message_type=SystemMessage.settings_changed, body=body)
|
|
VectorService(sensor).update_config()
|
|
|
|
def create(self, request, *args, **kwargs):
|
|
serializer = self.get_serializer(data=request.data)
|
|
serializer.is_valid(raise_exception=True)
|
|
sensor = serializer.save(type=DeviceType.SENSOR)
|
|
|
|
authorization_key = RabbitMQUserManagement(sensor.uuid).create_random_rabbitmq_credentials()
|
|
sensor.authorization_key = authorization_key
|
|
sensor.save()
|
|
|
|
VectorService(sensor).update_config()
|
|
|
|
return Response(SensorSerializer(instance=sensor).data, status=status.HTTP_201_CREATED)
|
|
|
|
def perform_destroy(self, sensor):
|
|
sensor.delete()
|
|
VectorService(sensor).delete_config()
|
|
|
|
@action(methods=['GET'], detail=True, url_path='system/(?P<system_action>[a-z]+)')
|
|
def system_message(self, request, system_action, *args, **kwargs):
|
|
sensor = self.get_object()
|
|
try:
|
|
action = SystemMessage[system_action]
|
|
except KeyError:
|
|
raise Http404
|
|
response = SensorManagement().send_message(sensor=sensor, message_type=action, wait_response=True)
|
|
return Response(response)
|
|
|
|
@action(methods=['GET'], detail=True, url_path='vector/(?P<vector_action>[a-z_]+)')
|
|
def vector_message(self, request, vector_action, *args, **kwargs):
|
|
sensor = self.get_object()
|
|
try:
|
|
action = VectorMessage[vector_action]
|
|
except KeyError:
|
|
raise Http404
|
|
response = SensorManagement().send_message(sensor=sensor, message_type=action, wait_response=True)
|
|
return Response(response)
|
|
|
|
@action(methods=['GET'], detail=True, url_path='zeek/(?P<zeek_action>[a-z_]+)')
|
|
def zeek_message(self, request, zeek_action, *args, **kwargs):
|
|
sensor = self.get_object()
|
|
try:
|
|
action = ZeekMessage[zeek_action]
|
|
except KeyError:
|
|
raise Http404
|
|
response = SensorManagement().send_message(sensor=sensor, message_type=action, wait_response=True)
|
|
|
|
return Response(response)
|
|
|
|
@action(methods=['GET'], detail=True, url_path='suricata/(?P<suricata_action>[a-z_]+)')
|
|
def suricata_message(self, request, suricata_action, *args, **kwargs):
|
|
sensor = self.get_object()
|
|
try:
|
|
action = SuricataMessage[suricata_action]
|
|
except KeyError:
|
|
raise Http404
|
|
response = SensorManagement().send_message(sensor=sensor, message_type=action, wait_response=True)
|
|
|
|
return Response(response)
|
|
|
|
@action(methods=['POST'], detail=True, url_path='zeek/protocols_disable')
|
|
def zeek_protocols_disable(self, request, *args, **kwargs):
|
|
serializer = ZeekProtocolsDisableSerializer(data=request.data)
|
|
sensor = self.get_object()
|
|
serializer.is_valid(raise_exception=True)
|
|
protocols = serializer.validated_data['disable_protocols']
|
|
response = SensorManagement().send_message(sensor=sensor,
|
|
message_type=ZeekMessage.protocols_disable,
|
|
wait_response=True,
|
|
body=protocols)
|
|
return Response(response)
|
|
|
|
@action(methods=['POST'], detail=True, url_path='zeek/settings_update')
|
|
def zeek_settings_update(self, request, *args, **kwargs):
|
|
sensor = self.get_object()
|
|
serializer = ZeekSettingsUpdateSerializer(data=request.data)
|
|
serializer.is_valid(raise_exception=True)
|
|
file = serializer.validated_data['file'].read()
|
|
response = SensorManagement().send_message(sensor=sensor,
|
|
message_type=ZeekMessage.settings_update,
|
|
wait_response=True,
|
|
body=file)
|
|
return Response(response)
|