From c20c63673454736141b1254759d94b46d3edb379 Mon Sep 17 00:00:00 2001
From: tdruez <tdruez@nexb.com>
Date: Mon, 30 Dec 2024 15:05:49 +0100
Subject: [PATCH] Add vulnerability_impact_notification DejacodeUser field #106

Signed-off-by: tdruez <tdruez@nexb.com>
---
 dje/admin.py                                  |  2 ++
 ...user_vulnerability_impact_notifications.py | 18 ++++++++++++++++++
 dje/models.py                                 | 14 ++++++++++++++
 vulnerabilities/fetch.py                      | 19 ++++++++++++++++---
 4 files changed, 50 insertions(+), 3 deletions(-)
 create mode 100644 dje/migrations/0006_dejacodeuser_vulnerability_impact_notifications.py

diff --git a/dje/admin.py b/dje/admin.py
index ef7ae47c..6393f89c 100644
--- a/dje/admin.py
+++ b/dje/admin.py
@@ -1347,6 +1347,7 @@ class DejacodeUserAdmin(
         "data_email_notification",
         "workflow_email_notification",
         "updates_email_notification",
+        "vulnerability_impact_notification",
         "company",
         "last_login",
         "last_api_access",
@@ -1379,6 +1380,7 @@ class DejacodeUserAdmin(
                     "data_email_notification",
                     "workflow_email_notification",
                     "updates_email_notification",
+                    "vulnerability_impact_notification",
                 )
             },
         ),
diff --git a/dje/migrations/0006_dejacodeuser_vulnerability_impact_notifications.py b/dje/migrations/0006_dejacodeuser_vulnerability_impact_notifications.py
new file mode 100644
index 00000000..3326921f
--- /dev/null
+++ b/dje/migrations/0006_dejacodeuser_vulnerability_impact_notifications.py
@@ -0,0 +1,18 @@
+# Generated by Django 5.0.9 on 2024-12-30 13:55
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('dje', '0005_dataspaceconfiguration_vulnerabilities_risk_threshold'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='dejacodeuser',
+            name='vulnerability_impact_notification',
+            field=models.BooleanField(default=False, help_text='Enable to receive internal notifications about new and updated vulnerabilities, including their impact on your packages and products.'),
+        ),
+    ]
diff --git a/dje/models.py b/dje/models.py
index be1a3448..00ae45de 100644
--- a/dje/models.py
+++ b/dje/models.py
@@ -1575,6 +1575,12 @@ def get_data_update_recipients(self, dataspace):
         # Need to be converted as a list to be serializable
         return list(qs.distinct().values_list("email", flat=True))
 
+    def get_vulnerability_notifications_users(self, dataspace):
+        """Return the Users with `vulnerability_notification` enabled for a given dataspace."""
+        return self.get_queryset().filter(
+            vulnerability_impact_notification=True, dataspace=dataspace
+        )
+
 
 class DejacodeUser(AbstractUser):
     uuid = models.UUIDField(
@@ -1616,6 +1622,14 @@ class DejacodeUser(AbstractUser):
         ),
     )
 
+    vulnerability_impact_notification = models.BooleanField(
+        default=False,
+        help_text=_(
+            "Enable to receive internal notifications about new and updated vulnerabilities, "
+            "including their impact on your packages and products."
+        ),
+    )
+
     company = models.CharField(
         max_length=30,
         blank=True,
diff --git a/vulnerabilities/fetch.py b/vulnerabilities/fetch.py
index d9eef65d..87d07a7d 100644
--- a/vulnerabilities/fetch.py
+++ b/vulnerabilities/fetch.py
@@ -8,11 +8,14 @@
 
 from timeit import default_timer as timer
 
+from django.contrib.auth import get_user_model
 from django.contrib.humanize.templatetags.humanize import intcomma
 from django.core.management.base import CommandError
 from django.urls import reverse
 from django.utils import timezone
 
+from notifications.signals import notify
+
 from component_catalog.models import PACKAGE_URL_FIELDS
 from component_catalog.models import Package
 from dejacode_toolkit.vulnerablecode import VulnerableCode
@@ -106,6 +109,7 @@ def fetch_for_packages(
 
         product_package_qs = ProductPackage.objects.filter(package__in=batch_affected_packages)
         product_package_qs.update_weighted_risk_score()
+        break
 
     return results
 
@@ -142,9 +146,9 @@ def notify_vulnerability_data_update(dataspace):
     VulnerableCode.
     """
     today = timezone.now().date()
-    vulnerability_qs = Vulnerability.objects.scope(dataspace).filter(last_modified_date__date=today)
+    vulnerability_qs = Vulnerability.objects.scope(dataspace)  #.filter(last_modified_date__date=today)
     package_qs = Package.objects.scope(dataspace).filter(
-        affected_by_vulnerabilities=vulnerability_qs
+        affected_by_vulnerabilities__in=vulnerability_qs
     )
     # product_qs = Product.objects.scope(dataspace).filter(packages=package_qs)
 
@@ -164,10 +168,19 @@ def notify_vulnerability_data_update(dataspace):
         f"{package_count} packages affected at {package_list_url}?is_vulnerable=yes\n"
     )
 
+    # 1. Webhooks
     find_and_fire_hook(
         "vulnerability.data_update",
         instance=None,
         dataspace=dataspace,
         payload_override={"text": f"{subject}\n{message}"},
     )
-    print(f"{subject}\n{message}")
+
+    # 2. Internal notifications
+    users_to_notify = get_user_model().objects.get_vulnerability_notifications_users(dataspace)
+    notify.send(
+        sender=Vulnerability,
+        verb="New vulnerabilities detected",
+        recipient=users_to_notify,
+        description=f"{message}",
+    )