Engine

api_app/analyzers_manager/migrations/0139_alter_analyzerconfig_mapping_data_model.py

@@ -6,13 +6,17 @@
 class Migration(migrations.Migration):
 
     dependencies = [
-        ('analyzers_manager', '0138_alter_analyzerreport_data_model_content_type'),
+        ("analyzers_manager", "0138_alter_analyzerreport_data_model_content_type"),
     ]
 
     operations = [
         migrations.AlterField(
-            model_name='analyzerconfig',
-            name='mapping_data_model',
-            field=models.JSONField(blank=True, default=dict, help_text='Mapping analyzer_report_key: data_model_key. Keys preceded by the symbol $ will be considered as constants.'),
+            model_name="analyzerconfig",
+            name="mapping_data_model",
+            field=models.JSONField(
+                blank=True,
+                default=dict,
+                help_text="Mapping analyzer_report_key: data_model_key. Keys preceded by the symbol $ will be considered as constants.",
+            ),
         ),
     ]

api_app/analyzers_manager/models.py

@@ -6,10 +6,8 @@
 
 from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
 from django.contrib.contenttypes.models import ContentType
-from django.contrib.postgres.fields import ArrayField
 from django.core.exceptions import ValidationError
 from django.db import models
-from django.db.models import ForeignKey
 
 from api_app.analyzers_manager.constants import (
     HashChoices,
@@ -19,12 +17,7 @@
 from api_app.analyzers_manager.exceptions import AnalyzerConfigurationException
 from api_app.analyzers_manager.queryset import AnalyzerReportQuerySet
 from api_app.choices import TLP, PythonModuleBasePaths
-from api_app.data_model_manager.models import (
-    BaseDataModel,
-    DomainDataModel,
-    FileDataModel,
-    IPDataModel,
-)
+from api_app.data_model_manager.models import BaseDataModel
 from api_app.fields import ChoiceArrayField
 from api_app.models import AbstractReport, PythonConfig, PythonModule
 
@@ -63,24 +56,9 @@ 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 +93,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 +107,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,
+            ),
         ),
     ]