diff --git a/api_app/analyzers_manager/models.py b/api_app/analyzers_manager/models.py index 6024d4c9b3..9052f99a6e 100644 --- a/api_app/analyzers_manager/models.py +++ b/api_app/analyzers_manager/models.py @@ -121,14 +121,20 @@ def _create_data_model_dictionary(self) -> Dict: return result def create_data_model(self) -> Optional[BaseDataModel]: - # TODO we don't need to actually crate a new object every time. - # if the report is the same of the previous one, we can just link it if not self._validation_before_data_model(): return None dictionary = self._create_data_model_dictionary() + temp_instance = self.data_model_class() + fingerprint = temp_instance.generate_fingerprint(data=dictionary) + existing_data_model = self.data_model_class.objects.filter(fingerprint=fingerprint).first() + + if existing_data_model: + logger.info(f"Deduplicated: Linking existing Data Model {existing_data_model.pk}") + self.data_model = existing_data_model + else: + self.data_model: BaseDataModel = self.data_model_class.objects.create(fingerprint=fingerprint) + self.data_model.merge(dictionary) - self.data_model: BaseDataModel = self.data_model_class.objects.create() - self.data_model.merge(dictionary) self.save() return self.data_model diff --git a/api_app/data_model_manager/migrations/0012_domaindatamodel_fingerprint_and_more.py b/api_app/data_model_manager/migrations/0012_domaindatamodel_fingerprint_and_more.py new file mode 100644 index 0000000000..33753b6589 --- /dev/null +++ b/api_app/data_model_manager/migrations/0012_domaindatamodel_fingerprint_and_more.py @@ -0,0 +1,25 @@ +from django.db import migrations, models + +class Migration(migrations.Migration): + + dependencies = [ + ('data_model_manager', '0011_data_model_date_index'), + ] + + operations = [ + migrations.AddField( + model_name='domaindatamodel', + name='fingerprint', + field=models.CharField(blank=True, db_index=True, max_length=64, default=''), + ), + migrations.AddField( + model_name='filedatamodel', + name='fingerprint', + field=models.CharField(blank=True, db_index=True, max_length=64, default=''), + ), + migrations.AddField( + model_name='ipdatamodel', + name='fingerprint', + field=models.CharField(blank=True, db_index=True, max_length=64, default=''), + ), + ] diff --git a/api_app/data_model_manager/migrations/0013_populate_fingerprints.py b/api_app/data_model_manager/migrations/0013_populate_fingerprints.py new file mode 100644 index 0000000000..6732de564d --- /dev/null +++ b/api_app/data_model_manager/migrations/0013_populate_fingerprints.py @@ -0,0 +1,87 @@ +import hashlib +import json +import logging + +from django.db import migrations + +logger = logging.getLogger(__name__) + +def normalize_dict(obj): + if isinstance(obj, dict): + return {k: normalize_dict(v) for k, v in sorted(obj.items())} + if isinstance(obj, list): + return [normalize_dict(i) for i in obj] + return obj + +def generate_fingerprint_from_instance(instance): + data = {} + for field in instance._meta.fields: + name = field.name + if name in ["id", "date", "fingerprint"]: + continue + value = getattr(instance, name) + if hasattr(value, "isoformat"): + value = value.isoformat() + data[name] = value + normalized_data = normalize_dict(data) + encoded_data = json.dumps(normalized_data, sort_keys=True).encode("utf-8") + return hashlib.sha256(encoded_data).hexdigest() + +def populate_fingerprints(apps, schema_editor): + batch_size = 500 + for model_name in ["IPDataModel", "DomainDataModel", "FileDataModel"]: + Model = apps.get_model("data_model_manager", model_name) + queryset = Model.objects.filter(fingerprint="").iterator(chunk_size=batch_size) + batch = [] + for instance in queryset: + try: + instance.fingerprint = generate_fingerprint_from_instance(instance) + batch.append(instance) + except Exception as e: + logger.error(f"Failed to generate fingerprint for {model_name} {instance.pk}: {e}") + if len(batch) >= batch_size: + Model.objects.bulk_update(batch, ["fingerprint"]) + batch = [] + if batch: + Model.objects.bulk_update(batch, ["fingerprint"]) + from django.contrib.contenttypes.models import ContentType + ct, _ = ContentType.objects.get_or_create(app_label="data_model_manager", model=model_name.lower()) + from django.db.models import Count + duplicates = Model.objects.values("fingerprint").annotate(c=Count("id")).filter(c__gt=1) + for entry in duplicates: + fp = entry["fingerprint"] + if not fp: + continue + instances = list(Model.objects.filter(fingerprint=fp).order_by("date")) + canonical = instances[0] + redundant_ids = [r.id for r in instances[1:]] + AnalyzerReport = apps.get_model("analyzers_manager", "AnalyzerReport") + AnalyzerReport.objects.filter( + data_model_content_type_id=ct.id, + data_model_object_id__in=redundant_ids + ).update(data_model_object_id=canonical.id) + Job = apps.get_model("api_app", "Job") + Job.objects.filter( + data_model_content_type_id=ct.id, + data_model_object_id__in=redundant_ids + ).update(data_model_object_id=canonical.id) + try: + UserAnalyzableEvent = apps.get_model("user_events_manager", "UserAnalyzableEvent") + UserAnalyzableEvent.objects.filter( + data_model_content_type_id=ct.id, + data_model_object_id__in=redundant_ids + ).update(data_model_object_id=canonical.id) + except LookupError: + pass + Model.objects.filter(id__in=redundant_ids).delete() + +def reverse_populate_fingerprints(apps, schema_editor): + pass + +class Migration(migrations.Migration): + dependencies = [ + ('data_model_manager', '0012_domaindatamodel_fingerprint_and_more'), + ] + operations = [ + migrations.RunPython(populate_fingerprints, reverse_populate_fingerprints), + ] diff --git a/api_app/data_model_manager/models.py b/api_app/data_model_manager/models.py index 4a2a04c5ef..e2cdd50682 100644 --- a/api_app/data_model_manager/models.py +++ b/api_app/data_model_manager/models.py @@ -1,6 +1,7 @@ +import hashlib import json import logging -from typing import Dict, Type, Union +from typing import Any, Dict, Type, Union from django.contrib.contenttypes.fields import GenericRelation from django.contrib.contenttypes.models import ContentType @@ -26,6 +27,14 @@ logger = logging.getLogger(__name__) +def normalize_dict(obj: Any) -> Any: + if isinstance(obj, dict): + return {k: normalize_dict(v) for k, v in sorted(obj.items())} + if isinstance(obj, list): + return [normalize_dict(i) for i in obj] + return obj + + class IETFReport(models.Model): rrname = LowercaseCharField(max_length=100) rrtype = LowercaseCharField(max_length=100) @@ -91,6 +100,7 @@ class BaseDataModel(models.Model): default=dict ) # field for additional information related to a specific analyzer date = models.DateTimeField(default=now) + fingerprint = models.CharField(max_length=64, db_index=True, blank=True, default="") analyzers_report = GenericRelation( to="analyzers_manager.AnalyzerReport", object_id_field="data_model_object_id", @@ -125,6 +135,30 @@ def owner(self) -> User: elif self.jobs.exists(): return self.jobs.first().user + def get_content_map(self, data: Dict = None) -> Dict: + if data is None: + data = {} + for field in self._meta.fields: + name = field.name + if name in ["id", "date", "fingerprint"]: + continue + value = getattr(self, name) + if hasattr(value, "isoformat"): + value = value.isoformat() + data[name] = value + data.pop("id", None) + data.pop("date", None) + data.pop("fingerprint", None) + data.pop("analyzers_report", None) + data.pop("jobs", None) + data.pop("user_events", None) + return normalize_dict(data) + + def generate_fingerprint(self, data: Dict = None) -> str: + content_map = self.get_content_map(data) + encoded_data = json.dumps(content_map, sort_keys=True).encode("utf-8") + return hashlib.sha256(encoded_data).hexdigest() + def merge(self, other: Union["BaseDataModel", Dict], append: bool = True) -> "BaseDataModel": if not self.pk: raise ValueError("Unable to merge a model that was not saved.") diff --git a/api_app/user_events_manager/signals.py b/api_app/user_events_manager/signals.py index e90cacd00d..22970bda7a 100644 --- a/api_app/user_events_manager/signals.py +++ b/api_app/user_events_manager/signals.py @@ -8,4 +8,5 @@ @receiver(models.signals.post_delete, sender=UserIPWildCardEvent) @receiver(models.signals.post_delete, sender=UserAnalyzableEvent) def post_delete_event_delete_data_model(sender, instance: UserDomainWildCardEvent, **kwargs): - instance.data_model.delete() + if instance.data_model: + instance.data_model.delete() diff --git a/tests/api_app/data_model_manager/test_cas_deduplication.py b/tests/api_app/data_model_manager/test_cas_deduplication.py new file mode 100644 index 0000000000..2a0ee5f2b7 --- /dev/null +++ b/tests/api_app/data_model_manager/test_cas_deduplication.py @@ -0,0 +1,71 @@ +from kombu import uuid + +from api_app.analyzables_manager.models import Analyzable +from api_app.analyzers_manager.models import AnalyzerConfig, AnalyzerReport +from api_app.choices import Classification +from api_app.data_model_manager.models import IPDataModel +from api_app.models import Job +from tests import CustomTestCase + + +class CASDeduplicationTestCase(CustomTestCase): + def setUp(self): + super().setUp() + self.analyzable = Analyzable.objects.get_or_create( + name="1.1.1.1", defaults={"classification": Classification.IP.value} + )[0] + self.job1 = Job.objects.create( + analyzable=self.analyzable, + status=Job.STATUSES.ANALYZERS_RUNNING.value, + ) + self.job2 = Job.objects.create( + analyzable=self.analyzable, + status=Job.STATUSES.ANALYZERS_RUNNING.value, + ) + self.config = AnalyzerConfig.objects.first() + + def test_smart_deduplication_via_fingerprint(self): + report1 = AnalyzerReport.objects.create( + job=self.job1, + config=self.config, + status=AnalyzerReport.STATUSES.SUCCESS.value, + task_id=str(uuid()), + parameters={}, + ) + report1._create_data_model_dictionary = lambda: {"isp": "Google", "asn": 15169} + dm1 = report1.create_data_model() + initial_count = IPDataModel.objects.count() + + report2 = AnalyzerReport.objects.create( + job=self.job2, + config=self.config, + status=AnalyzerReport.STATUSES.SUCCESS.value, + task_id=str(uuid()), + parameters={}, + ) + report2._create_data_model_dictionary = lambda: {"isp": "Google", "asn": 15169} + dm2 = report2.create_data_model() + self.assertEqual(dm1.pk, dm2.pk) + self.assertEqual(IPDataModel.objects.count(), initial_count) + + def test_normalization_stability(self): + report1 = AnalyzerReport.objects.create( + job=self.job1, + config=self.config, + status=AnalyzerReport.STATUSES.SUCCESS.value, + task_id=str(uuid()), + parameters={}, + ) + report1._create_data_model_dictionary = lambda: {"asn": 15169, "isp": "Google"} + dm1 = report1.create_data_model() + report2 = AnalyzerReport.objects.create( + job=self.job2, + config=self.config, + status=AnalyzerReport.STATUSES.SUCCESS.value, + task_id=str(uuid()), + parameters={}, + ) + report2._create_data_model_dictionary = lambda: {"isp": "Google", "asn": 15169} + dm2 = report2.create_data_model() + self.assertEqual(dm1.fingerprint, dm2.fingerprint) + self.assertEqual(dm1.pk, dm2.pk) diff --git a/tests/api_app/data_model_manager/test_migrations.py b/tests/api_app/data_model_manager/test_migrations.py new file mode 100644 index 0000000000..28decb4382 --- /dev/null +++ b/tests/api_app/data_model_manager/test_migrations.py @@ -0,0 +1,71 @@ +from django.db import connection +from django.db.migrations.executor import MigrationExecutor + +from api_app.helpers import calculate_md5, calculate_sha1, calculate_sha256 +from tests import CustomTestCase + + +class MigrationIntegrityTestCase(CustomTestCase): + @property + def app_name(self): + return "data_model_manager" + + @property + def migration_from(self): + return "0012_domaindatamodel_fingerprint_and_more" + + @property + def migration_to(self): + return "0013_populate_fingerprints" + + def setUp(self): + super().setUp() + self.executor = MigrationExecutor(connection) + self.old_state = self.executor.migrate([(self.app_name, self.migration_from)]) + + def test_migration_0013_deduplication_integrity(self): + old_apps = self.old_state.apps + IPDataModel = old_apps.get_model(self.app_name, "IPDataModel") + UserAnalyzableEvent = old_apps.get_model("user_events_manager", "UserAnalyzableEvent") + Analyzable = old_apps.get_model("analyzables_manager", "Analyzable") + ContentType = old_apps.get_model("contenttypes", "ContentType") + User = old_apps.get_model("certego_saas_user", "User") + user = User.objects.create(username="test_migrator", email="test@intelowl.org") + name1, name2 = "1.1.1.1", "8.8.8.8" + az1 = Analyzable.objects.create( + name=name1, + classification="ip", + md5=calculate_md5(name1.encode()), + sha1=calculate_sha1(name1.encode()), + sha256=calculate_sha256(name1.encode()), + ) + az2 = Analyzable.objects.create( + name=name2, + classification="ip", + md5=calculate_md5(name2.encode()), + sha1=calculate_sha1(name2.encode()), + sha256=calculate_sha256(name2.encode()), + ) + dm1 = IPDataModel.objects.create(evaluation="benign", reliability=5) + dm2 = IPDataModel.objects.create(evaluation="benign", reliability=5) + ct = ContentType.objects.get_for_model(IPDataModel) + UserAnalyzableEvent.objects.create( + user=user, analyzable=az1, data_model_content_type=ct, data_model_object_id=dm1.id + ) + UserAnalyzableEvent.objects.create( + user=user, analyzable=az2, data_model_content_type=ct, data_model_object_id=dm2.id + ) + self.executor.loader.build_graph() + new_state = self.executor.migrate([(self.app_name, self.migration_to)]) + new_apps = new_state.apps + IPDataModelNew = new_apps.get_model(self.app_name, "IPDataModel") + UserAnalyzableEventNew = new_apps.get_model("user_events_manager", "UserAnalyzableEvent") + self.assertEqual(IPDataModelNew.objects.count(), 1) + canonical = IPDataModelNew.objects.first() + events = UserAnalyzableEventNew.objects.filter(data_model_object_id=canonical.id) + self.assertEqual(events.count(), 2) + self.assertFalse(IPDataModelNew.objects.filter(id=dm2.id).exists()) + + def tearDown(self): + self.executor.migrate(self.executor.loader.graph.leaf_nodes()) + super().tearDown()