From 557146508616d930d7c0d4d4697f1e5e894d90f2 Mon Sep 17 00:00:00 2001
From: Renato Almeida de Oliveira <renato.almeida.oliveira@gmail.com>
Date: Sat, 7 Sep 2024 21:50:17 +0000
Subject: [PATCH] Add Signals to cascade VulnerabilityAssignment on Asset
 deletion, Fix #27

---
 develop/Dockerfile                            |  1 +
 nb_risk/__init__.py                           |  4 +++
 nb_risk/api/serializers.py                    |  4 +--
 nb_risk/models.py                             |  2 +-
 nb_risk/signals.py                            | 29 +++++++++++++++++++
 .../tests/vulnerabilityassignment/test_api.py |  2 +-
 6 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/develop/Dockerfile b/develop/Dockerfile
index 0a83423..efb97e9 100644
--- a/develop/Dockerfile
+++ b/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
diff --git a/nb_risk/__init__.py b/nb_risk/__init__.py
index 7cf2aad..69c4c35 100644
--- a/nb_risk/__init__.py
+++ b/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
diff --git a/nb_risk/api/serializers.py b/nb_risk/api/serializers.py
index 0200713..20d9bf6 100644
--- a/nb_risk/api/serializers.py
+++ b/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
 
diff --git a/nb_risk/models.py b/nb_risk/models.py
index 0e64de4..438a1af 100644
--- a/nb_risk/models.py
+++ b/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(
diff --git a/nb_risk/signals.py b/nb_risk/signals.py
index e69de29..59811ba 100644
--- a/nb_risk/signals.py
+++ b/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()
diff --git a/nb_risk/tests/vulnerabilityassignment/test_api.py b/nb_risk/tests/vulnerabilityassignment/test_api.py
index 13e45c6..527342a 100644
--- a/nb_risk/tests/vulnerabilityassignment/test_api.py
+++ b/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'   ]
\ No newline at end of file