Add Signals to cascade VulnerabilityAssignment on Asset deletion, Fix #…

…27

develop/Dockerfile

@@ -20,6 +20,7 @@ RUN git clone --single-branch --branch ${netbox_ver} https://github.com/netbox-c
 # Work around https://github.com/pallets/markupsafe/issues/284
 RUN pip install markupsafe==2.0.1
 RUN pip install git+https://github.com/python/tzdata
+RUN pip install tblib
 
 # -------------------------------------------------------------------------------------
 # Install Netbox Plugin

nb_risk/__init__.py

@@ -24,5 +24,9 @@ class NbriskConfig(PluginConfig):
         "additional_assets": [],
         "proxies": {},
     }
+
+    def ready(self):
+        from . import signals
+        super().ready()
 
 config = NbriskConfig  # noqa

nb_risk/api/serializers.py

@@ -102,7 +102,7 @@ class VulnerabilityAssignmentSerializer(NetBoxModelSerializer):
     asset = serializers.SerializerMethodField('get_asset')
     vulnerability = serializers.SlugRelatedField(slug_field="name", queryset=models.Vulnerability.objects.all())
 
-    asset_id = serializers.IntegerField(source='asset.id')
+    asset_id = serializers.IntegerField(source='asset.id', write_only=True)
 
     def validate(self, data):
         asset_id = data['asset']['id']
@@ -117,7 +117,7 @@ def get_asset(self, obj):
         serializer = get_serializer_for_model(obj.asset, prefix='Nested')
         context = {'request': self.context['request']}
         return serializer(obj.asset, context=context).data
-
+    
     def get_display(self, obj):
         return obj.name
 

nb_risk/models.py

@@ -116,7 +116,7 @@ def name(self):
 
     def __str__(self):
         return f"{self.asset} - {self.vulnerability.name}"
-
+   
     class Meta:
         constraints = (
             models.UniqueConstraint(

nb_risk/signals.py

@@ -0,0 +1,29 @@
+from django.db.models.signals import post_delete
+from django.dispatch import receiver
+from django.contrib.contenttypes.models import ContentType
+from netbox.plugins.utils import get_plugin_config
+
+
+from . import models
+
+@receiver(post_delete)
+def handle_vulnerable_asset_delete(sender,instance, **kwargs):
+
+    supported_model_class = False
+    supported_assets = get_plugin_config("nb_risk", "supported_assets")
+    additional_assets = get_plugin_config("nb_risk", "additional_assets")
+    supported_assets = supported_assets + additional_assets
+
+    for asset in supported_assets:
+        app_label, model = asset.split(".")
+        model = ContentType.objects.get(app_label=app_label, model=model).model_class()
+        if isinstance(instance, model):
+            supported_model_class = True
+            break
+
+    if supported_model_class:
+        related_VulAssings = models.VulnerabilityAssignment.objects.filter(
+            asset_object_type=ContentType.objects.get_for_model(instance),
+            asset_id=instance.id )
+        for vulnAssign in related_VulAssings:
+            vulnAssign.delete()

nb_risk/tests/vulnerabilityassignment/test_api.py

@@ -92,6 +92,6 @@ def setUpTestData(cls):
             },
         ]
 
-        cls.validation_excluded_fields = ['vulnerability']
+        cls.validation_excluded_fields = ['vulnerability', 'asset_id']
 
         cls.brief_fields = [ 'display', 'id', 'url'   ]