From c5c7a852be9a9a897486e05989243705f384168b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89ric=20Falconnier?= <eric@zentral.com>
Date: Mon, 25 Nov 2024 12:21:21 +0100
Subject: [PATCH] Fix Santa ballot box bug

The reset checks where not done per configuration, but globally.
---
 tests/santa/test_ballot_box.py      | 58 +++++++++++++++++++++++++++++
 zentral/contrib/santa/ballot_box.py | 11 +++++-
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/tests/santa/test_ballot_box.py b/tests/santa/test_ballot_box.py
index 115782565..2c4740d47 100644
--- a/tests/santa/test_ballot_box.py
+++ b/tests/santa/test_ballot_box.py
@@ -824,3 +824,61 @@ def test_ballot_box_target_state_reset(self):
         self.assertEqual(rules_qs.count(), 0)
         self.assertEqual(votes_qs.count(), 1)
         self.assertEqual(votes_qs.first(), vote)
+
+    # update target states
+
+    def test_update_target_states(self):
+        realm, realm_user = force_realm_user()
+        configuration = force_configuration(voting_realm=realm)
+        force_ballot(
+            self.file_target, realm_user,
+            [(configuration, True, configuration.partially_allowlisted_threshold)]
+        )
+        target_state, _ = TargetState.objects.update_or_create(
+            target=self.file_target,
+            configuration=configuration,
+            state=TargetState.State.UNTRUSTED,
+            reset_at=datetime.utcnow()
+        )
+        configuration2 = force_configuration(voting_realm=realm)
+        # second target state in unrelated configurations must not interfere
+        TargetState.objects.update_or_create(
+            target=self.file_target,
+            configuration=configuration2,
+            score=configuration2.partially_allowlisted_threshold,
+            state=TargetState.State.PARTIALLY_ALLOWLISTED,
+        )
+        ballot_box = BallotBox.for_realm_user(self.file_target, realm_user, all_configurations=True)
+        ballot_box._update_target_states([(configuration, True)])
+        target_state.refresh_from_db()
+        self.assertEqual(target_state.score, 0)
+
+    # partially allowlist
+
+    def test_partially_allowlist_rules(self):
+        realm, realm_user = force_realm_user()
+        configuration = force_configuration(voting_realm=realm)
+        force_ballot(
+            self.file_target, realm_user,
+            [(configuration, True, configuration.partially_allowlisted_threshold)]
+        )
+        target_state, _ = TargetState.objects.update_or_create(
+            target=self.file_target,
+            configuration=configuration,
+            state=TargetState.State.UNTRUSTED,
+            reset_at=datetime.utcnow()
+        )
+        self.assertEqual(configuration.rule_set.count(), 0)
+        configuration2 = force_configuration(voting_realm=realm)
+        # second target state in unrelated configurations must not interfere
+        TargetState.objects.update_or_create(
+            target=self.file_target,
+            configuration=configuration2,
+            score=configuration2.partially_allowlisted_threshold,
+            state=TargetState.State.PARTIALLY_ALLOWLISTED,
+        )
+        ballot_box = BallotBox.for_realm_user(self.file_target, realm_user, all_configurations=True)
+        with self.assertRaises(AssertionError) as cm:
+            ballot_box._partially_allowlist(configuration)
+        self.assertEqual(cm.exception.args[0], "No primary users found")
+        self.assertEqual(configuration.rule_set.count(), 0)
diff --git a/zentral/contrib/santa/ballot_box.py b/zentral/contrib/santa/ballot_box.py
index ded0c6412..165acd0ef 100644
--- a/zentral/contrib/santa/ballot_box.py
+++ b/zentral/contrib/santa/ballot_box.py
@@ -438,7 +438,10 @@ def _update_target_states(self, votes):
                     "select sum(v.weight * (case when v.was_yes_vote then 1 else -1 end)) "
                     "from santa_vote v "
                     "join santa_ballot b on (v.ballot_id = b.id) "
-                    "left join santa_targetstate ts on (b.target_id = ts.target_id) "
+                    "left join santa_targetstate ts on ("
+                    "  b.target_id = ts.target_id"
+                    "  and v.configuration_id = ts.configuration_id"
+                    ") "
                     "where b.target_id = %s and b.replaced_by_id is null "
                     "and (ts.reset_at is null or v.created_at > ts.reset_at) "
                     "and v.configuration_id = %s",
@@ -551,7 +554,10 @@ def _partially_allowlist(self, configuration):
                 "from santa_ballot b "
                 "left join realms_realmuser u on (b.realm_user_id = u.uuid) "
                 "join santa_vote v on (b.id = v.ballot_id) "
-                "join santa_targetstate ts on (b.target_id = ts.target_id) "
+                "join santa_targetstate ts on ("
+                "  b.target_id = ts.target_id"
+                "  and v.configuration_id = ts.configuration_id"
+                ") "
                 "where b.target_id = %s "
                 "and b.replaced_by_id is null "
                 "and v.configuration_id = %s "
@@ -560,6 +566,7 @@ def _partially_allowlist(self, configuration):
                 [self.target.pk, configuration.pk]
             )
             primary_users = list(r[0] for r in cursor.fetchall() if r)
+            assert len(primary_users) > 0, "No primary users found"
         self._update_or_create_rules(configuration, Rule.Policy.ALLOWLIST, primary_users)
 
     def _blocklist(self, configuration):