fixed migration ingestors (intelowlproject#2406)

federicofantini · Michalsus · commit a7d5261b9d27 · 2024-10-11T20:32:12.000+02:00
diff --git a/api_app/ingestors_manager/migrations/0021_ingestor_fix_malwarebazaar_threatfox.py b/api_app/ingestors_manager/migrations/0021_ingestor_fix_malwarebazaar_threatfox.py
@@ -1,33 +1,60 @@
 from django.db import migrations
+from django_celery_beat.models import PeriodicTask
 
 
 def migrate(apps, schema_editor):
     IngestorConfig = apps.get_model("ingestors_manager", "IngestorConfig")
 
     ic_threatfox = IngestorConfig.objects.get(name="ThreatFox")
     ic_malwarebazaar = IngestorConfig.objects.get(name="MalwareBazaar")
+    pt_threatfox = PeriodicTask.objects.get(name="ThreatFoxIngestor")
+    pt_malwarebazaar = PeriodicTask.objects.get(name="MalwareBazaarIngestor")
 
     ic_threatfox.disabled = True
     ic_malwarebazaar.disabled = True
+    pt_threatfox.enabled = False
+    pt_malwarebazaar.enabled = False
 
-    ic_threatfox.full_clean()
-    ic_threatfox.save()
+    ic_threatfox.routing_key = "ingestor"
+    ic_malwarebazaar.routing_key = "ingestor"
+
+    pt_threatfox.queue = "ingestor"
+    pt_malwarebazaar.queue = "ingestor"
 
+    ic_threatfox.full_clean()
     ic_malwarebazaar.full_clean()
+    pt_threatfox.full_clean()
+    pt_malwarebazaar.full_clean()
+
+    ic_threatfox.save()
     ic_malwarebazaar.save()
+    pt_threatfox.save()
+    pt_malwarebazaar.save()
 
 
 def reverse_migrate(apps, schema_editor):
     IngestorConfig = apps.get_model("ingestors_manager", "IngestorConfig")
 
     ic_threatfox = IngestorConfig.objects.get(name="ThreatFox")
     ic_malwarebazaar = IngestorConfig.objects.get(name="MalwareBazaar")
+    pt_threatfox = PeriodicTask.objects.get(name="ThreatFoxIngestor")
+    pt_malwarebazaar = PeriodicTask.objects.get(name="MalwareBazaarIngestor")
 
-    ic_threatfox.full_clean()
-    ic_threatfox.save()
+    ic_threatfox.routing_key = "default"
+    ic_malwarebazaar.routing_key = "default"
+
+    pt_threatfox.queue = "default"
+    pt_malwarebazaar.queue = "default"
 
+    ic_threatfox.full_clean()
     ic_malwarebazaar.full_clean()
+    pt_threatfox.full_clean()
+    pt_malwarebazaar.full_clean()
+
+    ic_threatfox.save()
     ic_malwarebazaar.save()
+    pt_threatfox.save()
+    pt_malwarebazaar.save()
 
 
 class Migration(migrations.Migration):
