Add Santa voting events

zentralopensource · Sep 17, 2024 · e8b4712 · e8b4712
1 parent 5e0c695
commit e8b4712
Show file tree

Hide file tree

Showing 11 changed files with 338 additions and 57 deletions.
diff --git a/server/realms/models.py b/server/realms/models.py
@@ -260,6 +260,26 @@ def groups_with_types(self):
             groups_with_types.append((realm_group, raw_groups[realm_group.pk]))
         return groups_with_types
 
+    def serialize_for_event(self, keys_only=False):
+        d = {
+            "pk": str(self.uuid),
+            "realm": self.realm.serialize_for_event(keys_only=True),
+            "username": self.username,
+        }
+        if keys_only:
+            return d
+        for attr in ("username",
+                     "email",
+                     "first_name",
+                     "last_name",
+                     "full_name",
+                     "custom_attr_1",
+                     "custom_attr_2"):
+            val = getattr(self, attr)
+            if val:
+                d[attr] = val
+        return d
+
 
 class RealmUserGroupMembership(models.Model):
     user = models.ForeignKey(RealmUser, on_delete=models.CASCADE)

diff --git a/tests/santa/test_ballots_views.py b/tests/santa/test_ballots_views.py
@@ -14,6 +14,7 @@
 from realms.backends.views import finalize_session
 from realms.models import RealmAuthenticationSession
 from zentral.contrib.santa.ballot_box import DuplicateVoteError, VotingNotAllowedError
+from zentral.contrib.santa.events import SantaBallotEvent, SantaTargetStateUpdateEvent
 from zentral.contrib.santa.models import Ballot, Target, TargetState
 from .utils import add_file_to_test_class, force_ballot, force_configuration, force_realm, force_realm_user
 
@@ -32,6 +33,7 @@ def setUpTestData(cls):
         cls.configuration = force_configuration(
             voting_realm=cls.realm,
             default_ballot_target_types=[Target.Type.METABUNDLE, Target.Type.BUNDLE, Target.Type.BINARY],
+            default_voting_weight=1,
         )
 
     # utility methods
@@ -387,18 +389,31 @@ def test_cast_ballot_post_voting_not_allowed(self, cast_votes):
         self.assertContains(response, "The ballot was rejected")
         self.assertEqual(ballot_qs.count(), 0)
 
-    def test_cast_ballot_post_yes(self):
+    @patch("zentral.core.queues.backends.kombu.EventQueues.post_event")
+    def test_cast_ballot_post_yes(self, post_event):
         ballot_qs = Ballot.objects.filter(target=self.metabundle_target)
         self.assertEqual(ballot_qs.count(), 0)
         self._login("santa.add_ballot", "santa.view_target", realm_user=True)
-        response = self.client.post(reverse("santa:cast_ballot")
-                                    + f"?target_type=METABUNDLE&target_identifier={self.metabundle_sha256}",
-                                    {f"cfg-{self.configuration.pk}-yes_no": "YES"},
-                                    follow=True)
+        with self.captureOnCommitCallbacks(execute=True) as callbacks:
+            response = self.client.post(reverse("santa:cast_ballot")
+                                        + f"?target_type=METABUNDLE&target_identifier={self.metabundle_sha256}",
+                                        {f"cfg-{self.configuration.pk}-yes_no": "YES"},
+                                        follow=True)
         self.assertEqual(response.status_code, 200)
+        self.assertEqual(len(callbacks), 1)
         self.assertTemplateUsed(response, "santa/target_detail.html")
         self.assertContains(response, "Your ballot has been cast")
         self.assertEqual(ballot_qs.count(), 1)
+        self.assertEqual(len(post_event.call_args_list), 3)
+        event1 = post_event.call_args_list[1].args[0]
+        self.assertIsInstance(event1, SantaBallotEvent)
+        self.assertEqual(len(event1.payload["votes"]), 1)
+        self.assertTrue(event1.payload["votes"][0]["was_yes_vote"])
+        self.assertEqual(event1.payload["votes"][0]["weight"], 1)
+        event2 = post_event.call_args_list[2].args[0]
+        self.assertIsInstance(event2, SantaTargetStateUpdateEvent)
+        self.assertEqual(event2.payload["new_value"]["score"], 1)
+        self.assertEqual(event2.payload["prev_value"]["score"], 0)
 
     def test_cast_ballot_post_no(self):
         ballot_qs = Ballot.objects.filter(target=self.file_target)

diff --git a/tests/santa/test_targets_views.py b/tests/santa/test_targets_views.py
@@ -895,16 +895,44 @@ def test_rest_target_state_post_not_allowed(self):
         ts.refresh_from_db()
         self.assertIsNone(ts.reset_at)
 
-    def test_rest_target_state_post_allowed(self):
+    @patch("zentral.core.queues.backends.kombu.EventQueues.post_event")
+    def test_rest_target_state_post_allowed(self, post_event):
         configuration = force_configuration()
         ts = TargetState.objects.create(configuration=configuration, target=self.file_target)
         self._login("santa.view_target", realm_user=True)
         force_voting_group(configuration, self.realm_user, can_reset_target=True)
-        response = self.client.post(reverse("santa:reset_target_state", args=(configuration.pk, ts.pk)),
-                                    follow=True)
+        with self.captureOnCommitCallbacks(execute=True) as callbacks:
+            response = self.client.post(reverse("santa:reset_target_state", args=(configuration.pk, ts.pk)),
+                                        follow=True)
         self.assertEqual(response.status_code, 200)
+        self.assertEqual(len(callbacks), 1)
         self.assertTemplateUsed(response, "santa/target_detail.html")
         self.assertContains(response, "Target state reset")
         self.assertNotContains(response, "Target state reset not allowed")
         ts.refresh_from_db()
         self.assertIsNotNone(ts.reset_at)
+        self.assertEqual(len(post_event.call_args_list), 2)
+        event = post_event.call_args_list[1].args[0]
+        self.assertEqual(
+            event.payload,
+            {'configuration': {'name': configuration.name, 'pk': configuration.pk},
+             'created_at': ts.created_at,
+             'new_value': {'flagged': False,
+                           'reset_at': ts.reset_at,
+                           'score': 0,
+                           'state': 0,
+                           'state_display': 'UNTRUSTED'},
+             'prev_value': {'flagged': False,
+                            'reset_at': None,
+                            'score': 0,
+                            'state': 0,
+                            'state_display': 'UNTRUSTED'},
+             'target': {'sha256': self.file_sha256,
+                        'type': 'BINARY'},
+             'updated_at': ts.updated_at}
+        )
+        self.assertEqual(
+            event.metadata.serialize()["objects"],
+            {'file': [f'sha256|{self.file_sha256}'],
+             'santa_configuration': [str(configuration.pk)]},
+        )
diff --git a/tests/santa/test_up_views.py b/tests/santa/test_up_views.py
@@ -1,5 +1,6 @@
 from datetime import datetime
 from importlib import import_module
+from unittest.mock import patch
 import uuid
 from django.conf import settings
 from django.http import HttpRequest
@@ -8,6 +9,7 @@
 from realms.backends.views import finalize_session
 from realms.models import RealmAuthenticationSession
 from zentral.contrib.santa.ballot_box import BallotBox
+from zentral.contrib.santa.events import SantaBallotEvent, SantaTargetStateUpdateEvent
 from zentral.contrib.santa.models import Ballot, Target
 from .utils import add_file_to_test_class, force_configuration, force_enrolled_machine, force_realm, force_realm_user
 
@@ -20,6 +22,7 @@ def setUpTestData(cls):
         cls.configuration = force_configuration(
             voting_realm=cls.realm,
             default_ballot_target_types=[Target.Type.METABUNDLE, Target.Type.SIGNING_ID],
+            default_voting_weight=1,
         )
         cls.em = force_enrolled_machine(
             configuration=cls.configuration,
@@ -227,14 +230,17 @@ def test_target_get_bundle_metabundle_ballot_box(self):
         self.assertEqual(ballot_box.target.type, Target.Type.METABUNDLE)
         self.assertEqual(ballot_box.target.identifier, self.metabundle_sha256)
 
-    def test_target_post_bundle_metabundle_ballot_box_yes(self):
+    @patch("zentral.core.queues.backends.kombu.EventQueues.post_event")
+    def test_target_post_bundle_metabundle_ballot_box_yes(self, post_event):
         self._login()
-        response = self.client.post(
-            reverse("realms_public:santa_up:target",
-                    args=(self.realm.pk, "bundle", self.bundle_sha256)),
-            {"yes_vote": "oui"},
-            follow=True
-        )
+        with self.captureOnCommitCallbacks(execute=True) as callbacks:
+            response = self.client.post(
+                reverse("realms_public:santa_up:target",
+                        args=(self.realm.pk, "bundle", self.bundle_sha256)),
+                {"yes_vote": "oui"},
+                follow=True
+            )
+        self.assertEqual(len(callbacks), 1)
         self.assertEqual(response.status_code, 200)
         self.assertTemplateUsed(response, "user_portal/santa_target_detail.html")
         self.assertContains(response, "Your ballot has been cast")
@@ -244,6 +250,42 @@ def test_target_post_bundle_metabundle_ballot_box_yes(self):
             response.context["existing_votes"],
             [(self.configuration, True)],
         )
+        self.assertEqual(len(post_event.call_args_list), 2)
+        event1 = post_event.call_args_list[0].args[0]
+        self.assertIsInstance(event1, SantaBallotEvent)
+        self.assertEqual(len(event1.payload["votes"]), 1)
+        self.assertEqual(event1.payload["votes"][0]["weight"], 1)
+        self.assertEqual(
+            event1.metadata.serialize()["objects"],
+            {"metabundle": [f'sha256|{self.metabundle_sha256}'],
+             "realm_user": [str(self.realm_user.pk)],
+             "santa_configuration": [str(self.configuration.pk)]}
+        )
+        event2 = post_event.call_args_list[1].args[0]
+        self.assertIsInstance(event2, SantaTargetStateUpdateEvent)
+        event2.payload.pop("created_at")
+        event2.payload.pop("updated_at")
+        self.assertEqual(
+            event2.payload,
+            {'configuration': {'name': self.configuration.name, 'pk': self.configuration.pk},
+             'new_value': {'flagged': False,
+                           'reset_at': None,
+                           'score': 1,
+                           'state': 0,
+                           'state_display': 'UNTRUSTED'},
+             'prev_value': {'flagged': False,
+                            'reset_at': None,
+                            'score': 0,
+                            'state': 0,
+                            'state_display': 'UNTRUSTED'},
+             'target': {'sha256': self.metabundle_sha256,
+                        'type': 'METABUNDLE'}}
+        )
+        self.assertEqual(
+            event2.metadata.serialize()["objects"],
+            {"metabundle": [f'sha256|{self.metabundle_sha256}'],
+             "santa_configuration": [str(self.configuration.pk)]}
+        )
 
     def test_target_post_bundle_metabundle_ballot_box_yolo(self):
         self._login()

diff --git a/tests/server_realms/test_realm_models.py b/tests/server_realms/test_realm_models.py
@@ -1,5 +1,5 @@
 from django.test import TestCase
-from .utils import force_realm, force_realm_group
+from .utils import force_realm, force_realm_group, force_realm_user
 
 
 class RealmModelsTestCase(TestCase):
@@ -64,3 +64,26 @@ def test_realm_group_serialize_for_event(self):
              'scim_external_id': None,
              'updated_at': realm_group.updated_at}
         )
+
+    def test_realm_user_serialize_for_event_keys_only(self):
+        realm, realm_user = force_realm_user()
+        self.assertEqual(
+            realm_user.serialize_for_event(keys_only=True),
+            {'pk': str(realm_user.pk),
+             'realm': {'name': realm.name,
+                       'pk': str(realm.pk)},
+             'username': realm_user.username}
+        )
+
+    def test_realm_user_serialize_for_event(self):
+        realm, realm_user = force_realm_user()
+        self.assertEqual(
+            realm_user.serialize_for_event(),
+            {'pk': str(realm_user.pk),
+             'realm': {'name': realm.name,
+                       'pk': str(realm.pk)},
+             'username': realm_user.username,
+             'email': realm_user.email,
+             'first_name': realm_user.first_name,
+             'last_name': realm_user.last_name}
+        )
diff --git a/zentral/contrib/santa/ballot_box.py b/zentral/contrib/santa/ballot_box.py
@@ -1,8 +1,11 @@
 from datetime import datetime
 import logging
+import uuid
 from django.db import connection
 from django.db.models import Q
 from django.utils.functional import cached_property
+from zentral.core.events.base import EventMetadata, EventRequest
+from .events import SantaBallotEvent, SantaTargetStateUpdateEvent
 from .models import Ballot, Configuration, EnrolledMachine, Rule, Target, TargetState, Vote, VotingGroup
 from .utils import target_related_targets
 
@@ -148,6 +151,7 @@ def __init__(self, target, voter, lock_target=True):
             self.target = target
         self.voter = voter
         self._set_target_states()
+        self._events = []
 
     def _set_target_states(self):
         self.target_states = {}
@@ -361,6 +365,7 @@ def _create_or_update_ballot(self, votes, event_target):
                 was_yes_vote=yes_vote,
                 weight=self.voter.voting_weight(configuration)
             )
+        self._events.append((SantaBallotEvent, ballot.serialize_for_event()))
         return ballot
 
     def _update_target_states(self, votes):
@@ -380,8 +385,22 @@ def _update_target_states(self, votes):
             new_score = result[0] or 0
             self._update_target_state(configuration, new_score, was_yes_vote)
 
+    def _queue_target_state_update_event(self, pre_update_state, target_state):
+        prev_value = {}
+        new_value = {}
+        post_update_state = target_state.serialize_for_event()
+        for attr in ("flagged", "reset_at", "score", "state", "state_display"):
+            prev_value[attr] = pre_update_state.get(attr)
+            new_value[attr] = post_update_state.pop(attr)
+        if prev_value != new_value:
+            event_payload = post_update_state
+            event_payload["new_value"] = new_value
+            event_payload["prev_value"] = prev_value
+            self._events.append((SantaTargetStateUpdateEvent, event_payload))
+
     def _update_target_state(self, configuration, score, was_yes_vote):
         target_state = self.target_states[configuration]
+        pre_update_state = target_state.serialize_for_event()
         if was_yes_vote:
             if target_state.flagged and self.voter.can_unflag_target(configuration):
                 target_state.flagged = False
@@ -394,6 +413,7 @@ def _update_target_state(self, configuration, score, was_yes_vote):
                 target_state.state = TargetState.State.SUSPECT
         target_state.score = score
         target_state.save()
+        self._queue_target_state_update_event(pre_update_state, target_state)
 
     def _update_target_state_state(self, target_state, score):
         configuration = target_state.configuration
@@ -423,11 +443,13 @@ def reset_target_state(self, configuration):
         if not self.voter.can_reset_target(configuration):
             raise ResetNotAllowedError
         target_state = self.target_states[configuration]
+        pre_update_state = target_state.serialize_for_event()
         target_state.score = 0
         target_state.flagged = False
         self._update_target_state_state(target_state, 0)
         target_state.reset_at = datetime.utcnow()
         target_state.save()
+        self._queue_target_state_update_event(pre_update_state, target_state)
 
     # rules
 
@@ -481,3 +503,17 @@ def _blocklist(self, configuration):
 
     def _ensure_no_rules(self, configuration):
         Rule.objects.filter(configuration=configuration, target__in=self._iter_rule_targets()).delete()
+
+    # events
+
+    def post_events(self, request, machine_serial_number=None):
+        event_request = EventRequest.build_from_request(request)
+        event_uuid = uuid.uuid4()
+        for index, (event_class, event_payload) in enumerate(self._events):
+            event_metadata = EventMetadata(
+                uuid=event_uuid, index=index,
+                request=event_request,
+                machine_serial_number=machine_serial_number,
+            )
+            event = event_class(event_metadata, event_payload)
+            event.post()