Job bi (#2052)

* job bi Signed-off-by: 0ssigeno <s.berni@certego.net> * More Signed-off-by: 0ssigeno <s.berni@certego.net> * Job bi Signed-off-by: 0ssigeno <s.berni@certego.net> * Migration Signed-off-by: 0ssigeno <s.berni@certego.net> * Fix Signed-off-by: 0ssigeno <s.berni@certego.net> * Fix Signed-off-by: 0ssigeno <s.berni@certego.net> * Added playbook in the elastic template Signed-off-by: 0ssigeno <s.berni@certego.net> * Fix Signed-off-by: 0ssigeno <s.berni@certego.net> * Missin migrations Signed-off-by: 0ssigeno <s.berni@certego.net> --------- Signed-off-by: 0ssigeno <s.berni@certego.net>
intelowlproject · Jan 9, 2024 · 8f68272 · 8f68272
1 parent 10583dd
commit 8f68272
Show file tree

Hide file tree

Showing 22 changed files with 283 additions and 78 deletions.
diff --git a/api_app/analyzers_manager/migrations/0057_analyzerreport_analyzerreportsbisearch.py b/api_app/analyzers_manager/migrations/0057_analyzerreport_analyzerreportsbisearch.py
@@ -0,0 +1,18 @@
+# Generated by Django 4.2.8 on 2024-01-08 15:26
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("analyzers_manager", "0056_alter_analyzer_config_dns0_rrsets_data"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="analyzerreport",
+            index=models.Index(
+                fields=["sent_to_bi", "-start_time"], name="analyzerreportsBISearch"
+            ),
+        ),
+    ]
diff --git a/api_app/analyzers_manager/models.py b/api_app/analyzers_manager/models.py
@@ -29,6 +29,7 @@ class AnalyzerReport(AbstractReport):
 
     class Meta:
         unique_together = [("config", "job")]
+        indexes = AbstractReport.Meta.indexes
 
 
 class MimeTypes(models.TextChoices):

diff --git a/api_app/analyzers_manager/queryset.py b/api_app/analyzers_manager/queryset.py
@@ -8,7 +8,7 @@
 
 class AnalyzerReportQuerySet(AbstractReportQuerySet):
     @classmethod
-    def _get_serializer_class(cls) -> Type["AnalyzerReportBISerializer"]:
+    def _get_bi_serializer_class(cls) -> Type["AnalyzerReportBISerializer"]:
         from api_app.analyzers_manager.serializers import AnalyzerReportBISerializer
 
         return AnalyzerReportBISerializer
diff --git a/api_app/connectors_manager/migrations/0028_connectorreport_connectorreportsbisearch.py b/api_app/connectors_manager/migrations/0028_connectorreport_connectorreportsbisearch.py
@@ -0,0 +1,18 @@
+# Generated by Django 4.2.8 on 2024-01-09 14:31
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("connectors_manager", "0027_connectorreport_sent_to_bi"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="connectorreport",
+            index=models.Index(
+                fields=["sent_to_bi", "-start_time"], name="connectorreportsBISearch"
+            ),
+        ),
+    ]
diff --git a/api_app/connectors_manager/models.py b/api_app/connectors_manager/models.py
@@ -17,6 +17,7 @@ class ConnectorReport(AbstractReport):
 
     class Meta:
         unique_together = [("config", "job")]
+        indexes = AbstractReport.Meta.indexes
 
 
 class ConnectorConfig(PythonConfig):

diff --git a/api_app/connectors_manager/queryset.py b/api_app/connectors_manager/queryset.py
@@ -8,7 +8,7 @@
 
 class ConnectorReportQuerySet(AbstractReportQuerySet):
     @classmethod
-    def _get_serializer_class(cls) -> Type["ConnectorReportBISerializer"]:
+    def _get_bi_serializer_class(cls) -> Type["ConnectorReportBISerializer"]:
         from api_app.connectors_manager.serializers import ConnectorReportBISerializer
 
         return ConnectorReportBISerializer
diff --git a/api_app/ingestors_manager/migrations/0015_ingestorreport_ingestorreportsbisearch.py b/api_app/ingestors_manager/migrations/0015_ingestorreport_ingestorreportsbisearch.py
@@ -0,0 +1,18 @@
+# Generated by Django 4.2.8 on 2024-01-09 14:31
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("ingestors_manager", "0014_ingestorreport_sent_to_bi"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="ingestorreport",
+            index=models.Index(
+                fields=["sent_to_bi", "-start_time"], name="ingestorreportsBISearch"
+            ),
+        ),
+    ]
diff --git a/api_app/ingestors_manager/models.py b/api_app/ingestors_manager/models.py
@@ -36,6 +36,7 @@ class IngestorReport(AbstractReport):
 
     class Meta:
         ordering = ["pk"]
+        indexes = AbstractReport.Meta.indexes
 
     def clean_report(self):
         if isinstance(self.report, list) and self.max_size_report is not None:

diff --git a/api_app/ingestors_manager/queryset.py b/api_app/ingestors_manager/queryset.py
@@ -8,7 +8,7 @@
 
 class IngestorReportQuerySet(AbstractReportQuerySet):
     @classmethod
-    def _get_serializer_class(cls) -> Type["IngestorReportBISerializer"]:
+    def _get_bi_serializer_class(cls) -> Type["IngestorReportBISerializer"]:
         from api_app.ingestors_manager.serializers import IngestorReportBISerializer
 
         return IngestorReportBISerializer
diff --git a/api_app/migrations/0053_job_sent_to_bi.py b/api_app/migrations/0053_job_sent_to_bi.py
@@ -0,0 +1,17 @@
+# Generated by Django 4.2.8 on 2024-01-08 14:13
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("api_app", "0052_periodic_task_bi"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="job",
+            name="sent_to_bi",
+            field=models.BooleanField(default=False, editable=False),
+        ),
+    ]
diff --git a/api_app/migrations/0054_job_jobbisearch.py b/api_app/migrations/0054_job_jobbisearch.py
@@ -0,0 +1,18 @@
+# Generated by Django 4.2.8 on 2024-01-08 15:24
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("api_app", "0053_job_sent_to_bi"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="job",
+            index=models.Index(
+                fields=["sent_to_bi", "-received_request_time"], name="JobBISearch"
+            ),
+        ),
+    ]
diff --git a/api_app/models.py b/api_app/models.py
@@ -219,6 +219,9 @@ class Meta:
                 fields=["playbook_to_execute", "finished_analysis_time", "user"],
                 name="PlaybookConfigOrdering",
             ),
+            models.Index(
+                fields=["sent_to_bi", "-received_request_time"], name="JobBISearch"
+            ),
         ]
 
     # constants
@@ -306,6 +309,7 @@ class Meta:
     scan_check_time = models.DurationField(
         null=True, blank=True, default=datetime.timedelta(hours=24)
     )
+    sent_to_bi = models.BooleanField(editable=False, default=False)
 
     def __str__(self):
         return f'{self.__class__.__name__}(#{self.pk}, "{self.analyzed_object_name}")'
@@ -914,6 +918,11 @@ class AbstractReport(models.Model):
 
     class Meta:
         abstract = True
+        indexes = [
+            models.Index(
+                fields=["sent_to_bi", "-start_time"], name="%(class)ssBISearch"
+            )
+        ]
 
     def __str__(self):
         return f"{self.__class__.__name__}(job:#{self.job_id}, {self.config.name})"

diff --git a/api_app/pivots_manager/migrations/0022_pivotreport_pivotreportsbisearch.py b/api_app/pivots_manager/migrations/0022_pivotreport_pivotreportsbisearch.py
@@ -0,0 +1,18 @@
+# Generated by Django 4.2.8 on 2024-01-09 14:31
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("pivots_manager", "0021_pivotreport_sent_to_bi"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="pivotreport",
+            index=models.Index(
+                fields=["sent_to_bi", "-start_time"], name="pivotreportsBISearch"
+            ),
+        ),
+    ]
diff --git a/api_app/pivots_manager/models.py b/api_app/pivots_manager/models.py
@@ -28,6 +28,7 @@ class PivotReport(AbstractReport):
 
     class Meta:
         unique_together = [("config", "job")]
+        indexes = AbstractReport.Meta.indexes
 
 
 class PivotMap(models.Model):

diff --git a/api_app/pivots_manager/queryset.py b/api_app/pivots_manager/queryset.py
@@ -36,7 +36,7 @@ def valid(
 
 class PivotReportQuerySet(AbstractReportQuerySet):
     @classmethod
-    def _get_serializer_class(cls) -> Type["PivotReportBISerializer"]:
+    def _get_bi_serializer_class(cls) -> Type["PivotReportBISerializer"]:
         from api_app.pivots_manager.serializers import PivotReportBISerializer
 
         return PivotReportBISerializer
diff --git a/api_app/queryset.py b/api_app/queryset.py
@@ -10,7 +10,7 @@
 
 if TYPE_CHECKING:
     from api_app.models import PythonConfig
-    from api_app.serializers import AbstractReportBISerializer
+    from api_app.serializers import AbstractBIInterface
 
 from celery.canvas import Signature
 from django.db import models
@@ -37,6 +37,57 @@
 from certego_saas.apps.user.models import User
 
 
+class SendToBiQuerySet(models.QuerySet):
+    @classmethod
+    def _get_bi_serializer_class(cls) -> Type["AbstractBIInterface"]:
+        raise NotImplementedError()
+
+    @staticmethod
+    def _create_index_template():
+        if not settings.ELASTICSEARCH_CLIENT.indices.exists_template(
+            name=settings.ELASTICSEARCH_BI_INDEX
+        ):
+            with open(
+                settings.CONFIG_ROOT / "elastic_search_mappings" / "intel_owl_bi.json"
+            ) as f:
+                body = json.load(f)
+                body["index_patterns"] = [f"{settings.ELASTICSEARCH_BI_INDEX}-*"]
+                settings.ELASTICSEARCH_CLIENT.indices.put_template(
+                    name=settings.ELASTICSEARCH_BI_INDEX, body=body
+                )
+
+    def send_to_elastic_as_bi(self, max_timeout: int = 60) -> bool:
+        from elasticsearch.helpers import bulk
+
+        BULK_MAX_SIZE = 1000
+        found_errors = False
+
+        p = Paginator(self, BULK_MAX_SIZE)
+        for i in p.page_range:
+            page = p.get_page(i)
+            objects = page.object_list
+            serializer = self._get_bi_serializer_class()(instance=objects, many=True)
+            objects_serialized = serializer.data
+            _, errors = bulk(
+                settings.ELASTICSEARCH_CLIENT,
+                objects_serialized,
+                request_timeout=max_timeout,
+            )
+            if errors:
+                logging.error(
+                    f"Errors on sending to elastic: {errors}."
+                    " We are not marking objects as sent."
+                )
+                found_errors |= errors
+            else:
+                logging.info("BI sent")
+                self.model.objects.filter(
+                    pk__in=objects.values_list("pk", flat=True)
+                ).update(sent_to_bi=True)
+        self._create_index_template()
+        return found_errors
+
+
 class CleanOnCreateQuerySet(models.QuerySet):
     def create(self, **kwargs):
         obj = self.model(**kwargs)
@@ -78,7 +129,13 @@ def annotate_runnable(self, user: User = None) -> QuerySet:
         return self.annotate(runnable=Exists(qs))
 
 
-class JobQuerySet(CleanOnCreateQuerySet):
+class JobQuerySet(CleanOnCreateQuerySet, SendToBiQuerySet):
+    @classmethod
+    def _get_bi_serializer_class(cls):
+        from api_app.serializers import JobBISerializer
+
+        return JobBISerializer
+
     def visible_for_user(self, user: User) -> "JobQuerySet":
         """
         User has access to:
@@ -230,55 +287,8 @@ def annotate_value_for_user(
         )
 
 
-class AbstractReportQuerySet(QuerySet):
-    @classmethod
-    def _get_serializer_class(cls) -> Type["AbstractReportBISerializer"]:
-        raise NotImplementedError()
-
-    @staticmethod
-    def _create_index_template():
-        if not settings.ELASTICSEARCH_CLIENT.indices.exists_template(
-            name=settings.ELASTICSEARCH_BI_INDEX
-        ):
-            with open(
-                settings.CONFIG_ROOT / "elastic_search_mappings" / "intel_owl_bi.json"
-            ) as f:
-                body = json.load(f)
-                body["index_patterns"] = [f"{settings.ELASTICSEARCH_BI_INDEX}-*"]
-                settings.ELASTICSEARCH_CLIENT.indices.put_template(
-                    name=settings.ELASTICSEARCH_BI_INDEX, body=body
-                )
-
-    def send_to_elastic_as_bi(self, max_timeout: int = 60) -> bool:
-        from elasticsearch.helpers import bulk
-
-        BULK_MAX_SIZE = 1000
-        found_errors = False
-
-        p = Paginator(self.order_by("pk"), BULK_MAX_SIZE)
-        for i in p.page_range:
-            page = p.get_page(i)
-            objects: AbstractReportQuerySet = page.object_list
-            serializer = self._get_serializer_class()(instance=objects, many=True)
-            objects_serialized = serializer.data
-            _, errors = bulk(
-                settings.ELASTICSEARCH_CLIENT,
-                objects_serialized,
-                request_timeout=max_timeout,
-            )
-            if errors:
-                logging.error(
-                    f"Errors on sending to elastic: {errors}."
-                    " We are not marking objects as sent."
-                )
-                found_errors |= errors
-            else:
-                logging.info("BI sent")
-                self.model.objects.filter(
-                    pk__in=objects.values_list("pk", flat=True)
-                ).update(sent_to_bi=True)
-        self._create_index_template()
-        return found_errors
+class AbstractReportQuerySet(SendToBiQuerySet):
+    ...
 
 
 class ModelWithOwnershipQuerySet: