Add Santa targets counters

zentral/contrib/santa/events/__init__.py

@@ -3,6 +3,7 @@
 from zentral.conf import settings
 from zentral.contrib.inventory.models import File
 from zentral.contrib.santa.models import Bundle, EnrolledMachine, Target
+from zentral.contrib.santa.utils import add_bundle_binary_targets, update_or_create_targets
 from zentral.core.events.base import BaseEvent, EventMetadata, EventRequest, register_event_type
 from zentral.utils.certificates import APPLE_DEV_ID_ISSUER_CN, parse_apple_dev_id
 from zentral.utils.text import shard
@@ -263,6 +264,16 @@ def _build_file_tree_from_santa_event(event_d):
     return app_d
 
 
+def _is_allow_event(event_d):
+    decision = event_d.get('decision')
+    return decision and decision.startswith("ALLOW_")
+
+
+def _is_block_event(event_d):
+    decision = event_d.get('decision')
+    return decision and decision.startswith("BLOCK_")
+
+
 def _is_allow_unknown_event(event_d):
     return event_d.get('decision') == "ALLOW_UNKNOWN"
 
@@ -271,7 +282,56 @@ def _is_bundle_binary_pseudo_event(event_d):
     return event_d.get('decision') == "BUNDLE_BINARY"
 
 
-def _create_missing_bundles(events):
+def _update_targets(events):
+    targets = {}
+    for event_d in events:
+        # target keys
+        target_keys = []
+        file_sha256 = event_d.get("file_sha256")
+        if file_sha256:
+            target_keys.append((Target.BINARY, file_sha256))
+        cdhash = event_d.get("cdhash")
+        if cdhash:
+            target_keys.append((Target.CDHASH, cdhash))
+        team_id = event_d.get("team_id")
+        if team_id:
+            target_keys.append((Target.TEAM_ID, team_id))
+        signing_id = event_d.get("signing_id")
+        if signing_id:
+            target_keys.append((Target.SIGNING_ID, signing_id))
+        signing_chain = event_d.get("signing_chain")
+        if signing_chain:
+            for cert_d in signing_chain:
+                sha256 = cert_d.get("sha256")
+                if sha256:
+                    target_keys.append((Target.CERTIFICATE, sha256))
+        if not _is_bundle_binary_pseudo_event(event_d):
+            bundle_hash = event_d.get("file_bundle_hash")
+            if bundle_hash:
+                target_keys.append((Target.BUNDLE, bundle_hash))
+        # increments
+        blocked_incr = collected_incr = executed_incr = 0
+        if _is_block_event(event_d):
+            blocked_incr = 1
+        elif _is_bundle_binary_pseudo_event(event_d):
+            collected_incr = 1
+        elif _is_allow_event(event_d):
+            executed_incr = 1
+        else:
+            logger.warning("Unknown decision")
+        # aggregations
+        for target_key in target_keys:
+            target_increments = targets.setdefault(
+                target_key,
+                {"blocked_incr": 0, "collected_incr": 0, "executed_incr": 0}
+            )
+            target_increments["blocked_incr"] += blocked_incr
+            target_increments["collected_incr"] += collected_incr
+            target_increments["executed_incr"] += executed_incr
+    return update_or_create_targets(targets)
+
+
+def _create_missing_bundles(events, targets):
     bundle_events = {
         sha256: event_d
         for sha256, event_d in (
@@ -292,7 +352,10 @@ def _create_missing_bundles(events):
     )
     unknown_file_bundle_hashes = list(set(bundle_events.keys()) - existing_sha256_set)
     for sha256 in unknown_file_bundle_hashes:
-        target, _ = Target.objects.get_or_create(type=Target.BUNDLE, identifier=sha256)
+        target, _ = targets.get((Target.BUNDLE, sha256), (None, None))
+        if not target:
+            logger.error("Missing BUNDLE target %s", sha256)
+            continue
         defaults = {}
         event_d = bundle_events[sha256]
         for event_attr, bundle_attr in (("file_bundle_path", "path"),
@@ -329,17 +392,15 @@ def _create_bundle_binaries(events):
         if bundle.uploaded_at:
             logger.info("Bundle %s already uploaded", bundle_sha256)
             continue
-        binary_targets = []
+        binary_target_identifiers = []
         binary_count = bundle.binary_count
         for event_d in events:
             if not binary_count:
                 event_binary_count = event_d.get("file_bundle_binary_count")
                 if event_binary_count:
                     binary_count = event_binary_count
-            binary_sha256 = event_d.get("file_sha256")
-            binary_target, _ = Target.objects.get_or_create(type=Target.BINARY, identifier=binary_sha256)
-            binary_targets.append(binary_target)
-        bundle.binary_targets.add(*binary_targets)
+            binary_target_identifiers.append(event_d["file_sha256"])
+        add_bundle_binary_targets(bundle, binary_target_identifiers)
         save_bundle = False
         if not bundle.binary_count and binary_count:
             bundle.binary_count = binary_count
@@ -411,7 +472,8 @@ def process_events(enrolled_machine, user_agent, ip, data):
     events = data.get("events", [])
     if not events:
         return []
-    unknown_file_bundle_hashes = _create_missing_bundles(events)
+    targets = _update_targets(events)
+    unknown_file_bundle_hashes = _create_missing_bundles(events, targets)
     _create_bundle_binaries(events)
     _commit_files(events)
     _post_santa_events(enrolled_machine, user_agent, ip, events)

zentral/contrib/santa/migrations/0035_target_blocked_count_target_collected_count_and_more.py

@@ -0,0 +1,40 @@
+# Generated by Django 4.2.13 on 2024-06-08 10:03
+
+from django.db import migrations, models
+import django.utils.timezone
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('santa', '0034_enrolledmachine_cdhash_rule_count_alter_target_type'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='target',
+            name='blocked_count',
+            field=models.IntegerField(default=0),
+        ),
+        migrations.AddField(
+            model_name='target',
+            name='collected_count',
+            field=models.IntegerField(default=0),
+        ),
+        migrations.AddField(
+            model_name='target',
+            name='created_at',
+            field=models.DateTimeField(auto_now_add=True, default=django.utils.timezone.now),
+            preserve_default=False,
+        ),
+        migrations.AddField(
+            model_name='target',
+            name='executed_count',
+            field=models.IntegerField(default=0),
+        ),
+        migrations.AddField(
+            model_name='target',
+            name='updated_at',
+            field=models.DateTimeField(auto_now=True),
+        ),
+    ]

zentral/contrib/santa/models.py

@@ -660,6 +660,11 @@ class Target(models.Model):
     )
     type = models.CharField(choices=TYPE_CHOICES, max_length=16)
     identifier = models.CharField(max_length=64)
+    blocked_count = models.IntegerField(default=0)
+    collected_count = models.IntegerField(default=0)
+    executed_count = models.IntegerField(default=0)
+    created_at = models.DateTimeField(auto_now_add=True)
+    updated_at = models.DateTimeField(auto_now=True)
 
     objects = TargetManager()
 

zentral/contrib/santa/utils.py

@@ -1,7 +1,11 @@
+from datetime import datetime
 import plistlib
 from dateutil import parser
+from django.db import connection
+import psycopg2.extras
 from zentral.conf import settings
 from zentral.utils.payloads import generate_payload_uuid, get_payload_identifier, sign_payload
+from .models import Target
 
 
 def build_santa_enrollment_configuration(enrollment):
@@ -100,3 +104,46 @@ def parse_santa_log_message(message):
     if args:
         d["args"] = args.split()
     return d
+
+
+def update_or_create_targets(targets):
+    query = (
+        'insert into santa_target '
+        '("type", "identifier", "blocked_count", "collected_count", "executed_count", "created_at", "updated_at") '
+        'values %s '
+        'on conflict ("type", "identifier") do update '
+        'set blocked_count = santa_target.blocked_count + excluded.blocked_count, '
+        'collected_count = santa_target.collected_count + excluded.collected_count, '
+        'executed_count = santa_target.executed_count + excluded.executed_count, '
+        'updated_at = excluded.updated_at '
+        'returning *, (xmax = 0) AS _created'
+    )
+    with connection.cursor() as cursor:
+        now = datetime.utcnow()
+        result = psycopg2.extras.execute_values(
+            cursor, query,
+            ((target_type, target_identifier,
+              val["blocked_incr"], val["collected_incr"], val["executed_incr"],
+              now, now)
+             for (target_type, target_identifier), val in targets.items()),
+            fetch=True
+        )
+        columns = [c.name for c in cursor.description]
+        targets = {}
+        for t in result:
+            target_d = dict(zip(columns, t))
+            created = target_d.pop("_created")
+            target = Target(**target_d)
+            targets[(target.type, target.identifier)] = (target, created)
+        return targets
+
+
+def add_bundle_binary_targets(bundle, binary_target_identifiers):
+    query = (
+        'insert into santa_bundle_binary_targets '
+        '("bundle_id", "target_id") '
+        "select %s, id from santa_target where type = 'BINARY' and identifier in %s "
+        'on conflict ("bundle_id", "target_id") do nothing'
+    )
+    with connection.cursor() as cursor:
+        return cursor.execute(query, [bundle.pk, tuple(binary_target_identifiers)])