Engine + Analyzable (#2685)

* engine

* More stuff

* More engine

* More engine

* Removed old code

* Blake

* Fixes

* Fixes

* Fixes

* Fixes

* Fix

* Fix merge with dict

* Fix

* More tests

* Added another engine

* Analyzable

* Blake

* Fixes

* Fixes

* Fixes

* Added error

* Added delete

* Fixes

* Fixes

* Blake

* Fix

* More fixes

* Fixes

* Fixes

* update

* Fixes

* Fixes

* Fix typo

* More fixes

* More fixes

* Fixed files

* Update deepsource

* Fixes

* Typo

* Fixes

* Fixes

* Fixes

* Fixes

* Blake

* Fix

* Fix

* Typo

Author: 0ssigeno

api_app/admin.py

@@ -49,10 +49,8 @@ class JobAdminView(CustomAdminView):
         "id",
         "status",
         "user",
-        "observable_name",
-        "observable_classification",
-        "file_name",
-        "file_mimetype",
+        "get_analyzable_name",
+        "get_analyzable_classification",
         "received_request_time",
         "analyzers_executed",
         "connectors_executed",
@@ -64,13 +62,16 @@ class JobAdminView(CustomAdminView):
         "user",
         "status",
     )
-    search_fields = (
-        "md5",
-        "observable_name",
-        "file_name",
-    )
     list_filter = ("status", "user", "tags")
 
+    @admin.display(description="Name")
+    def get_analyzable_name(self, instance):
+        return instance.analyzable.name
+
+    @admin.display(description="Classification")
+    def get_analyzable_classification(self, instance):
+        return instance.analyzable.classification
+
     @staticmethod
     def has_add_permission(request: HttpRequest) -> bool:
         return False

api_app/analyzables_manager/__init__.py

@@ -0,0 +1,11 @@
+from django.contrib import admin
+
+from api_app.analyzables_manager.models import Analyzable
+
+
+@admin.register(Analyzable)
+class AnalyzableAdmin(admin.ModelAdmin):
+    list_display = ["pk", "name", "sha1", "sha256", "md5"]
+    search_fields = ["name", "sha1", "sha256", "md5"]
+    ordering = ["name"]
+    list_filter = ["discovery_date"]

api_app/analyzables_manager/admin.py

@@ -0,0 +1,70 @@
+# Generated by Django 4.2.17 on 2025-01-22 08:59
+
+from django.db import migrations, models
+from django.utils.timezone import now
+
+import api_app.defaults
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = []
+
+    operations = [
+        migrations.CreateModel(
+            name="Analyzable",
+            fields=[
+                (
+                    "id",
+                    models.BigAutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                ("md5", models.CharField(max_length=255, unique=True, editable=False)),
+                (
+                    "sha256",
+                    models.CharField(max_length=255, unique=True, editable=False),
+                ),
+                ("sha1", models.CharField(max_length=255, unique=True, editable=False)),
+                ("name", models.CharField(max_length=255)),
+                (
+                    "mimetype",
+                    models.CharField(
+                        blank=True, max_length=80, null=True, default=None
+                    ),
+                ),
+                (
+                    "file",
+                    models.FileField(
+                        null=True,
+                        default=None,
+                        blank=True,
+                        upload_to=api_app.defaults.file_directory_path,
+                    ),
+                ),
+                (
+                    "classification",
+                    models.CharField(
+                        max_length=100,
+                        choices=[
+                            ("ip", "Ip"),
+                            ("url", "Url"),
+                            ("domain", "Domain"),
+                            ("hash", "Hash"),
+                            ("generic", "Generic"),
+                            ("file", "File"),
+                        ],
+                    ),
+                ),
+                ("discovery_date", models.DateTimeField(default=now)),
+            ],
+            options={
+                "abstract": False,
+            },
+        ),
+    ]

api_app/analyzables_manager/apps.py

@@ -0,0 +1,78 @@
+import hashlib
+
+from django.db import migrations
+from django.db.models import F, OuterRef, Subquery, Window
+from django.db.models.functions import RowNumber
+
+
+def calculate_sha1(value: bytes) -> str:
+    return hashlib.sha1(value).hexdigest()  # skipcq BAN-B324
+
+
+def calculate_sha256(value: bytes) -> str:
+    return hashlib.sha256(value).hexdigest()  # skipcq BAN-B324
+
+
+def migrate(apps, schema_editor):
+    Job = apps.get_model("api_app", "Job")
+    Analyzable = apps.get_model("analyzables_manager", "Analyzable")
+    # get only on job for md5
+    jobs = Job.objects.alias(
+        row_number=Window(
+            RowNumber(), partition_by=(F("md5"),), order_by="received_request_time"
+        )
+    ).filter(row_number=1)
+    for job in jobs:
+        if job.is_sample:
+            an = Analyzable.objects.create(
+                md5=job.md5,
+                sha256=job.sha256,
+                sha1=job.sha1,
+                file=job.file,
+                mimetype=job.file_mimetype,
+                name=job.file_name,
+                classification="sample",
+                discovery_date=job.received_request_time,
+            )
+
+            p = job.file.path
+            try:
+                p.rename(p.parent / job.md5)
+            except Exception:
+                print(f"Error: unable to rename {job}")
+            else:
+                job.file.name = job.md5
+            with open(p, "rb") as f:
+                content = f.read()
+                f.seek(0)
+            an.sha1 = calculate_sha1(content)
+            an.sha256 = calculate_sha256(content)
+        else:
+            an = Analyzable.objects.create(
+                md5=job.md5,
+                name=job.observable_name,
+                classification=job.observable_classification,
+                discovery_date=job.received_request_time,
+            )
+            an.sha1 = calculate_sha1(an.name.encode("utf-8"))
+            an.sha256 = calculate_sha256(an.name.encode("utf-8"))
+        an.save()
+    Job.objects.update(
+        analyzable=Subquery(
+            Analyzable.objects.filter(md5=OuterRef("md5")).values("pk")[:1]
+        )
+    )
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("contenttypes", "0002_remove_content_type_name"),
+        ("api_app", "0067_add_analyzable"),
+        ("analyzables_manager", "0001_initial"),
+        ("visualizers_manager", "0040_visualizer_config_data_model"),
+    ]
+
+    operations = [
+        migrations.RunPython(migrate, migrations.RunPython.noop),
+    ]

api_app/analyzables_manager/migrations/0001_initial.py

@@ -0,0 +1,25 @@
+# Generated by Django 4.2.17 on 2025-01-23 14:38
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("analyzables_manager", "0002_migrate_data"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="analyzable",
+            index=models.Index(
+                fields=["classification"], name="analyzables_classif_adf7ca_idx"
+            ),
+        ),
+        migrations.AddIndex(
+            model_name="analyzable",
+            index=models.Index(
+                fields=["mimetype"], name="analyzables_mimetyp_321d7d_idx"
+            ),
+        ),
+    ]

api_app/analyzables_manager/migrations/0002_migrate_data.py

@@ -0,0 +1,96 @@
+from typing import Type, Union
+
+from django.core.exceptions import ValidationError
+from django.db import models
+from django.utils.timezone import now
+
+from api_app.analyzables_manager.queryset import AnalyzableQuerySet
+from api_app.choices import Classification
+from api_app.data_model_manager.models import (
+    BaseDataModel,
+    DomainDataModel,
+    FileDataModel,
+    IPDataModel,
+)
+from api_app.defaults import file_directory_path
+from api_app.helpers import calculate_md5, calculate_sha1, calculate_sha256
+
+
+class Analyzable(models.Model):
+    name = models.CharField(max_length=255)
+    discovery_date = models.DateTimeField(default=now)
+    md5 = models.CharField(max_length=255, unique=True, editable=False)
+    sha256 = models.CharField(max_length=255, unique=True, editable=False)
+    sha1 = models.CharField(max_length=255, unique=True, editable=False)
+    classification = models.CharField(max_length=100, choices=Classification.choices)
+    mimetype = models.CharField(max_length=80, blank=True, null=True, default=None)
+    file = models.FileField(
+        upload_to=file_directory_path, null=True, blank=True, default=None
+    )
+
+    objects = AnalyzableQuerySet.as_manager()
+
+    class Meta:
+        indexes = [
+            models.Index(fields=["classification"]),
+            models.Index(fields=["mimetype"]),
+        ]
+
+    def __str__(self):
+        return self.name
+
+    @property
+    def analyzed_object(self):
+        return self.file if self.is_sample else self.name
+
+    @property
+    def is_sample(self) -> bool:
+        return self.classification == Classification.FILE.value
+
+    def get_data_model_class(self) -> Type[BaseDataModel]:
+        if self.classification == Classification.IP.value:
+            return IPDataModel
+        elif self.classification in [
+            Classification.URL.value,
+            Classification.DOMAIN.value,
+        ]:
+            return DomainDataModel
+        elif self.classification in [
+            Classification.HASH.value,
+            Classification.FILE.value,
+        ]:
+            return FileDataModel
+        else:
+            raise NotImplementedError()
+
+    def _set_hashes(self, value: Union[str, bytes]):
+        if isinstance(value, str):
+            value = value.encode("utf-8")
+        if not self.md5:
+            self.md5 = calculate_md5(value)
+        if not self.sha256:
+            self.sha256 = calculate_sha256(value)
+        if not self.sha1:
+            self.sha1 = calculate_sha1(value)
+
+    def clean(self):
+        if self.classification == Classification.FILE.value:
+            from api_app.analyzers_manager.models import MimeTypes
+
+            if not self.file:
+                raise ValidationError("File must be set for samples")
+            content = self.read()
+            if not self.mimetype:
+                self.mimetype = MimeTypes.calculate(content, self.name)
+        else:
+            if self.mimetype or self.file:
+                raise ValidationError(
+                    "Mimetype and file must not be set for observables"
+                )
+            content = self.name
+        self._set_hashes(content)
+
+    def read(self) -> bytes:
+        if self.classification == Classification.FILE.value:
+            self.file.seek(0)
+            return self.file.read()

api_app/analyzables_manager/migrations/0003_analyzable_analyzables_classif_adf7ca_idx_and_more.py

@@ -0,0 +1,31 @@
+import logging
+
+from django.db.models import QuerySet
+
+logger = logging.getLogger(__name__)
+
+
+class AnalyzableQuerySet(QuerySet):
+
+    def visible_for_user(self, user):
+
+        from api_app.models import Job
+
+        analyzables = (
+            Job.objects.visible_for_user(user)
+            .values("analyzable")
+            .distinct()
+            .values_list("analyzable__pk", flat=True)
+        )
+        return self.filter(pk__in=analyzables)
+
+    def create(self, *args, **kwargs):
+        obj = self.model(**kwargs)
+        self._for_write = True
+        try:
+            obj.full_clean()
+        except Exception as e:
+            logger.error(f"Already exists obj {obj.md5}")
+            raise e
+        obj.save(force_insert=True, using=self.db)
+        return obj