feat: add scan scheduling options to data source model

Introduced `ScanScheduling` enum to distinguish between manual and scheduled scans. Enforced restrictions on scan uploads and queueing based on the scan scheduling type. Includes a migration to add the `scan_scheduling` field to the `DataSource` model.
CentreForDigitalHumanities · Jan 13, 2025 · 78c8a9b · 78c8a9b
1 parent 8d13bf7
commit 78c8a9b
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 4 deletions.
diff --git a/humitifier-server/src/api/views/system.py b/humitifier-server/src/api/views/system.py
@@ -8,7 +8,7 @@
 
 from api.permissions import TokenHasApplication
 from api.serializers import DataSourceSyncSerializer
-from hosts.models import DataSource, DataSourceType, Host
+from hosts.models import DataSource, DataSourceType, Host, ScanScheduling
 
 
 class UploadScans(APIView):
@@ -41,7 +41,14 @@ def post(self, request, format=None):
             if "host" not in scan or "data" not in scan:
                 raise APIException("Invalid scan data")
 
-            host, created = Host.objects.get_or_create(fqdn=scan["host"])
+            host, created = Host.objects.get(fqdn=scan["host"])
+
+            if (
+                host.data_source
+                and host.data_source.scan_scheduling != ScanScheduling.MANUAL
+            ):
+                return Response("Cannot upload scan for non-manually scheduled hosts", status=status.HTTP_400_BAD_REQUEST)
+
             host.add_scan(scan["data"])
 
             hosts.append(host)

diff --git a/humitifier-server/src/hosts/management/commands/queue_scan.py b/humitifier-server/src/hosts/management/commands/queue_scan.py
@@ -1,6 +1,6 @@
-from django.core.management.base import BaseCommand
+from django.core.management.base import BaseCommand, CommandError
 
-from hosts.models import Host
+from hosts.models import Host, ScanScheduling
 
 
 class Command(BaseCommand):
@@ -13,4 +13,12 @@ def handle(self, *args, **options):
 
         host = Host.objects.get(fqdn=options["host"])
 
+        if (
+            host.data_source
+            and host.data_source.scan_scheduling != ScanScheduling.SCHEDULED
+        ):
+            raise CommandError(
+                "Cannot queue scans for hosts using non-scheduled scheduling"
+            )
+
         start_scan.delay(host.fqdn)
diff --git a/humitifier-server/src/hosts/migrations/0017_datasource_scan_scheduling.py b/humitifier-server/src/hosts/migrations/0017_datasource_scan_scheduling.py
@@ -0,0 +1,25 @@
+# Generated by Django 5.1.4 on 2025-01-13 10:02
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("hosts", "0016_host_otap_stage"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="datasource",
+            name="scan_scheduling",
+            field=models.CharField(
+                choices=[
+                    ("manual", "Manual/host initiated scanning"),
+                    ("scheduled", "Scheduled scanning"),
+                ],
+                default="scheduled",
+                max_length=255,
+            ),
+        ),
+    ]
diff --git a/humitifier-server/src/hosts/models.py b/humitifier-server/src/hosts/models.py
@@ -46,6 +46,11 @@ class DataSourceType(models.TextChoices):
     MANUAL = ("manual", "Manual entry")
 
 
+class ScanScheduling(models.TextChoices):
+    MANUAL = ("manual", "Manual/host initiated scanning")
+    SCHEDULED = ("scheduled", "Scheduled scanning")
+
+
 class AlertLevel(models.TextChoices):
     INFO = "info"
     WARNING = "warning"
@@ -123,6 +128,10 @@ class Meta:
         max_length=255, choices=DataSourceType.choices, default=DataSourceType.MANUAL
     )
 
+    scan_scheduling = models.CharField(
+        max_length=255, choices=ScanScheduling.choices, default=ScanScheduling.SCHEDULED
+    )
+
     def __str__(self):
         return self.name