import logging from django.conf import settings from django.contrib.auth.models import User from django.db import models from django.db.models.signals import post_save from django.dispatch import receiver from django.utils.translation import gettext_lazy, pgettext_lazy from django_celery_beat.models import PeriodicTask from rest_framework.authtoken.models import Token from solo.models import SingletonModel from core.fields import IntegerField _log = logging.getLogger(__name__) @receiver(post_save, sender=settings.AUTH_USER_MODEL) def create_auth_token(sender, instance=None, created=False, **kwargs): """ Generate Auth token for user See https://www.django-rest-framework.org/api-guide/authentication/#generating-tokens """ if created: Token.objects.create(user=instance) class NameDescriptionModel(models.Model): """ Abstract model for name - description pair model """ name = models.CharField(max_length=128, verbose_name=gettext_lazy( 'Name')) description = models.TextField(null=True, blank=True, verbose_name=gettext_lazy('Description'), help_text=gettext_lazy('Description')) class Meta: abstract = True ordering = ['name'] def __str__(self): return self.name class UniqueNameDescriptionModel(NameDescriptionModel): """ NameDescriptionModel heir to make name field unique """ name = models.CharField(max_length=128, verbose_name=gettext_lazy('Name'), unique=True) class Meta: abstract = True ordering = ['name'] class SensorConnectedMixin(models.Model): """ Add connection to sensor by name. Sensor can be ARMAIF or Endpoint """ sensor = models.CharField(null=True, blank=True, max_length=128, verbose_name=gettext_lazy("Sensor name")) class Meta: abstract = True class UpdatedNameDescriptionModel(NameDescriptionModel): """ Abstract model for name - description pair with updated info """ updated = models.DateTimeField(auto_now=True, verbose_name=gettext_lazy( 'Updated'), help_text=gettext_lazy('Date and time, when asset was updated')) class Meta: abstract = True ordering = ['updated', 'name'] def __str__(self): return self.name class VulnerabilityEffect(NameDescriptionModel): """ Possible Vulnerability effect """ pass class VulnerabilityRecommendations(NameDescriptionModel): """ Recommendations how to close Vulnerability """ pass class Vulnerability(UniqueNameDescriptionModel): """ Description of one Vulnerability """ detection_method = models.TextField(verbose_name=gettext_lazy('Detection method')) affected_software = models.TextField(verbose_name=gettext_lazy('Vulnerable software')) close_recommendations = models.ManyToManyField(VulnerabilityRecommendations, verbose_name=gettext_lazy('Resolve recommendations'), help_text=gettext_lazy('How to resolve the vulnerability'), blank=True) effects = models.ManyToManyField(VulnerabilityEffect, verbose_name=pgettext_lazy('as consequences', 'Effects'), help_text=gettext_lazy('Consequences of the vulnerability'), blank=True) class ConnectionType(NameDescriptionModel): """ Asset connection type """ pass class Connection(models.Model): """ Connection between two assets """ from assets.models.assets import Asset class ProtocolType(models.TextChoices): TCP = 'TCP', gettext_lazy('TCP protocol') UDP = 'UDP', gettext_lazy('UDP protocol') src_asset = models.ForeignKey(Asset, related_name='src', on_delete=models.CASCADE, verbose_name=gettext_lazy('Source asset')) dst_asset = models.ForeignKey(Asset, related_name='dst', on_delete=models.CASCADE, verbose_name=gettext_lazy('Destination asset')) src_port = IntegerField(verbose_name=gettext_lazy('Source port'), null=True, blank=True, min_value=0, max_value=65535) dst_port = IntegerField(verbose_name=gettext_lazy('Destination port'), null=True, blank=True, min_value=0, max_value=65535) connection_protocol = models.CharField(choices=ProtocolType.choices, verbose_name=gettext_lazy('Connection protocol'), help_text=gettext_lazy('Connection protocol type'), blank=True, null=True, default=ProtocolType.TCP, max_length=10) created = models.DateTimeField(auto_now=True) updated = models.DateTimeField(auto_now=True, verbose_name=gettext_lazy( 'Updated'), help_text=gettext_lazy('Date and time, when connection was updated'))