Engine

api_app/analyzers_manager/models.py

@@ -63,24 +63,11 @@ def clean(self):
         ):
             raise ValidationError("Wrong data model for this report")
 
-    @classmethod
-    def get_data_model_class(cls, job) -> Type[BaseDataModel]:
-        if job.is_sample or job.observable_classification == ObservableTypes.HASH.value:
-            return FileDataModel
-        if job.observable_classification == ObservableTypes.IP.value:
-            return IPDataModel
-        if job.observable_classification in [
-            ObservableTypes.DOMAIN.value,
-            ObservableTypes.URL.value,
-        ]:
-            return DomainDataModel
-        raise NotImplementedError(
-            f"Unable to find data model for {job.observable_classification}"
-        )
+
 
     @property
     def data_model_class(self) -> Type[BaseDataModel]:
-        return self.get_data_model_class(self.job)
+        return self.job.get_data_model_class()
 
     def _validation_before_data_model(self) -> bool:
         if not self.status == self.STATUSES.SUCCESS.value:
@@ -115,7 +102,6 @@ def _create_data_model_dictionary(self) -> Dict:
         result = {"malware_family": "MalwareFamily"}.
         """
         result = {}
-        data_model_fields = self.data_model_class.get_fields()
         logger.debug(f"Mapping is {json.dumps(self.config.mapping_data_model)}")
         for report_key, data_model_key in self.config.mapping_data_model.items():
             # this is a constant
@@ -130,40 +116,18 @@ def _create_data_model_dictionary(self) -> Dict:
                     # validation
                     self.errors.append(f"Field {report_key} not available in report")
                     continue
-
-            # create the related object if necessary
-            if isinstance(data_model_fields[data_model_key], ForeignKey):
-                # to create an object we need at least a dictionary
-                if not isinstance(value, dict):
-                    self.errors.append(
-                        f"Field {report_key} has type {type(report_key)} while a dictionary is expected"
-                    )
-                    continue
-                value, _ = data_model_fields[
-                    data_model_key
-                ].related_model.objects.get_or_create(**value)
-                result[data_model_key] = value
-            elif isinstance(data_model_fields[data_model_key], ArrayField):
-                if data_model_key not in result:
-                    result[data_model_key] = []
-                if isinstance(value, list):
-                    result[data_model_key].extend(value)
-                elif isinstance(value, dict):
-                    result[data_model_key].extend(list(value.keys()))
-                else:
-                    result[data_model_key].append(value)
-            else:
-                result[data_model_key] = value
+            result[data_model_key] = value
         return result
 
     def create_data_model(self) -> Optional[BaseDataModel]:
         if not self._validation_before_data_model():
             return None
         dictionary = self._create_data_model_dictionary()
-        data_model = self.data_model_class.objects.create(**dictionary)
-        self.data_model = data_model
+
+        self.data_model: BaseDataModel = self.data_model_class.objects.create()
+        self.data_model.merge(dictionary)
         self.save()
-        return data_model
+        return self.data_model
 
 
 class MimeTypes(models.TextChoices):

api_app/analyzers_manager/observable_analyzers/crowdsec.py

@@ -61,14 +61,14 @@ def _update_data_model(self, data_model):
                     self.report.data_model_class.EVALUATIONS.CLEAN.value
                 )
             elif "Proxy" in label or "VPN" in label:
-                data_model.tags = [DataModelTags.ANONYMIZER]
+                data_model.tags = [DataModelTags.ANONYMIZER.value]
                 data_model.evaluation = (
                     self.report.data_model_class.EVALUATIONS.CLEAN.value
                 )
             elif label in ["TOR exit node"]:
                 data_model.tags = [
-                    DataModelTags.ANONYMIZER,
-                    DataModelTags.TOR_EXIT_NODE,
+                    DataModelTags.ANONYMIZER.value,
+                    DataModelTags.TOR_EXIT_NODE.value,
                 ]
                 data_model.evaluation = (
                     self.report.data_model_class.EVALUATIONS.CLEAN.value

api_app/analyzers_manager/queryset.py

@@ -16,7 +16,7 @@ def _get_bi_serializer_class(cls) -> Type["AnalyzerReportBISerializer"]:
         return AnalyzerReportBISerializer
 
     def get_data_models(self, job) -> QuerySet:
-        DataModel = self.model.get_data_model_class(job)  # noqa
+        DataModel = job.get_data_model_class()  # noqa
         return DataModel.objects.filter(
             pk__in=self.values_list("data_model_object_id", flat=True)
         )

api_app/data_model_manager/migrations/0005_alter_domaindatamodel_external_references_and_more.py

@@ -8,68 +8,134 @@
 class Migration(migrations.Migration):
 
     dependencies = [
-        ('data_model_manager', '0004_alter_domaindatamodel_evaluation_and_more'),
+        ("data_model_manager", "0004_alter_domaindatamodel_evaluation_and_more"),
     ]
 
     operations = [
         migrations.AlterField(
-            model_name='domaindatamodel',
-            name='external_references',
-            field=api_app.data_model_manager.fields.SetField(base_field=models.URLField(), blank=True, default=list, size=None),
+            model_name="domaindatamodel",
+            name="external_references",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=models.URLField(), blank=True, default=list, size=None
+            ),
         ),
         migrations.AlterField(
-            model_name='domaindatamodel',
-            name='related_threats',
-            field=api_app.data_model_manager.fields.SetField(base_field=api_app.data_model_manager.fields.LowercaseCharField(max_length=100), blank=True, default=list, size=None),
+            model_name="domaindatamodel",
+            name="related_threats",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=api_app.data_model_manager.fields.LowercaseCharField(
+                    max_length=100
+                ),
+                blank=True,
+                default=list,
+                size=None,
+            ),
         ),
         migrations.AlterField(
-            model_name='domaindatamodel',
-            name='resolutions',
-            field=api_app.data_model_manager.fields.SetField(base_field=api_app.data_model_manager.fields.LowercaseCharField(max_length=100), default=list, size=None),
+            model_name="domaindatamodel",
+            name="resolutions",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=api_app.data_model_manager.fields.LowercaseCharField(
+                    max_length=100
+                ),
+                default=list,
+                size=None,
+            ),
         ),
         migrations.AlterField(
-            model_name='domaindatamodel',
-            name='tags',
-            field=api_app.data_model_manager.fields.SetField(base_field=api_app.data_model_manager.fields.LowercaseCharField(max_length=100), blank=True, default=None, null=True, size=None),
+            model_name="domaindatamodel",
+            name="tags",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=api_app.data_model_manager.fields.LowercaseCharField(
+                    max_length=100
+                ),
+                blank=True,
+                default=None,
+                null=True,
+                size=None,
+            ),
         ),
         migrations.AlterField(
-            model_name='filedatamodel',
-            name='comments',
-            field=api_app.data_model_manager.fields.SetField(base_field=api_app.data_model_manager.fields.LowercaseCharField(max_length=100), blank=True, default=list, size=None),
+            model_name="filedatamodel",
+            name="comments",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=api_app.data_model_manager.fields.LowercaseCharField(
+                    max_length=100
+                ),
+                blank=True,
+                default=list,
+                size=None,
+            ),
         ),
         migrations.AlterField(
-            model_name='filedatamodel',
-            name='external_references',
-            field=api_app.data_model_manager.fields.SetField(base_field=models.URLField(), blank=True, default=list, size=None),
+            model_name="filedatamodel",
+            name="external_references",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=models.URLField(), blank=True, default=list, size=None
+            ),
         ),
         migrations.AlterField(
-            model_name='filedatamodel',
-            name='related_threats',
-            field=api_app.data_model_manager.fields.SetField(base_field=api_app.data_model_manager.fields.LowercaseCharField(max_length=100), blank=True, default=list, size=None),
+            model_name="filedatamodel",
+            name="related_threats",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=api_app.data_model_manager.fields.LowercaseCharField(
+                    max_length=100
+                ),
+                blank=True,
+                default=list,
+                size=None,
+            ),
         ),
         migrations.AlterField(
-            model_name='filedatamodel',
-            name='tags',
-            field=api_app.data_model_manager.fields.SetField(base_field=api_app.data_model_manager.fields.LowercaseCharField(max_length=100), blank=True, default=None, null=True, size=None),
+            model_name="filedatamodel",
+            name="tags",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=api_app.data_model_manager.fields.LowercaseCharField(
+                    max_length=100
+                ),
+                blank=True,
+                default=None,
+                null=True,
+                size=None,
+            ),
         ),
         migrations.AlterField(
-            model_name='ipdatamodel',
-            name='external_references',
-            field=api_app.data_model_manager.fields.SetField(base_field=models.URLField(), blank=True, default=list, size=None),
+            model_name="ipdatamodel",
+            name="external_references",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=models.URLField(), blank=True, default=list, size=None
+            ),
         ),
         migrations.AlterField(
-            model_name='ipdatamodel',
-            name='related_threats',
-            field=api_app.data_model_manager.fields.SetField(base_field=api_app.data_model_manager.fields.LowercaseCharField(max_length=100), blank=True, default=list, size=None),
+            model_name="ipdatamodel",
+            name="related_threats",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=api_app.data_model_manager.fields.LowercaseCharField(
+                    max_length=100
+                ),
+                blank=True,
+                default=list,
+                size=None,
+            ),
         ),
         migrations.AlterField(
-            model_name='ipdatamodel',
-            name='resolutions',
-            field=api_app.data_model_manager.fields.SetField(base_field=models.URLField(), default=list, size=None),
+            model_name="ipdatamodel",
+            name="resolutions",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=models.URLField(), default=list, size=None
+            ),
         ),
         migrations.AlterField(
-            model_name='ipdatamodel',
-            name='tags',
-            field=api_app.data_model_manager.fields.SetField(base_field=api_app.data_model_manager.fields.LowercaseCharField(max_length=100), blank=True, default=None, null=True, size=None),
+            model_name="ipdatamodel",
+            name="tags",
+            field=api_app.data_model_manager.fields.SetField(
+                base_field=api_app.data_model_manager.fields.LowercaseCharField(
+                    max_length=100
+                ),
+                blank=True,
+                default=None,
+                null=True,
+                size=None,
+            ),
         ),
     ]

api_app/data_model_manager/models.py

@@ -1,10 +1,13 @@
 import json
-from typing import Dict, Type
+from typing import Dict, Type, Union
 
 from django.contrib.contenttypes.fields import GenericRelation
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.postgres import fields as pg_fields
+from django.contrib.postgres.fields import ArrayField
 from django.db import models
+from django.db.models import ManyToManyField, ForeignKey
+from django.forms import JSONField
 from django.utils.timezone import now
 from rest_framework.serializers import ModelSerializer
 
@@ -16,7 +19,9 @@
 from api_app.data_model_manager.fields import LowercaseCharField, SetField
 from api_app.data_model_manager.queryset import BaseDataModelQuerySet
 from certego_saas.apps.user.models import User
+import logging
 
+logger = logging.getLogger(__name__)
 
 class IETFReport(models.Model):
     rrname = LowercaseCharField(max_length=100)
@@ -103,6 +108,11 @@ class BaseDataModel(models.Model):
         object_id_field="data_model_object_id",
         content_type_field="data_model_content_type",
     )
+    job = GenericRelation(
+        to="api_app.Job",
+        object_id_field="data_model_object_id",
+        content_type_field="data_model_content_type",
+    )
 
     TAGS = DataModelTags
 
@@ -111,6 +121,46 @@ class BaseDataModel(models.Model):
     class Meta:
         abstract = True
 
+    def merge(self, other: Union["BaseDataModel", Dict], append:bool=True) -> "BaseDataModel":
+        assert self.pk
+        if not isinstance(other, (self.__class__, dict)):
+            raise TypeError(f"Different class between {self} and {type(other)}")
+        for field_name, field in self.get_fields().items():
+            if field_name == "id":
+                continue
+            result_attr = getattr(self, field_name)
+            if isinstance(other, dict):
+                other_attr = other[field_name]
+            else:
+                other_attr = getattr(other, field_name, None)
+            if not other_attr:
+                continue
+            if append:
+                if isinstance(field, ArrayField):
+                    result_attr.extend(other_attr)
+                elif isinstance(field, (JSONField, SetField)):
+                    result_attr |= other_attr
+                elif isinstance(field, ManyToManyField):
+                    result_attr.add(*other_attr.values_list("pk", flat=True))
+                    continue
+                elif isinstance(field, ForeignKey):
+                    if isinstance(other_attr, dict):
+                        other_attr = field.related_model.objects.get_or_create(**other_attr)
+                    elif isinstance(other_attr, models.Model):
+                        pass
+                    else:
+                        logger.error(f"Field {field_name} has wrong type with value {other_attr}")
+                        continue
+                    result_attr = other_attr
+                else:
+                    result_attr = other_attr
+            else:
+                result_attr = other_attr
+            setattr(self, field_name, result_attr)
+        self.save()
+        return self
+
+
     @classmethod
     def get_content_type(cls) -> ContentType:
         return ContentType.objects.get_for_model(model=cls)
@@ -129,6 +179,9 @@ def owner(self) -> User:
     def get_serializer(cls) -> Type[ModelSerializer]:
         raise NotImplementedError()
 
+    def serialize(self) -> Dict:
+        return self.get_serializer()(self, read_only=True).data
+
 
 class DomainDataModel(BaseDataModel):
     ietf_report = models.ManyToManyField(IETFReport, related_name="domains")  # pdns