From 8109c9c4255bb01296e058e16c88cbe7c396cc50 Mon Sep 17 00:00:00 2001
From: James Campbell <james.campbell@tanti.org.uk>
Date: Thu, 8 Aug 2024 14:57:26 +0200
Subject: [PATCH] Add infraction type enum for granular punishment

---
 .../contracts/coordination/InfractionCollector.sol    | 11 ++++++++---
 tests/test_infraction.py                              |  6 +-----
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/contracts/contracts/coordination/InfractionCollector.sol b/contracts/contracts/coordination/InfractionCollector.sol
index f66af146a..c4992068e 100644
--- a/contracts/contracts/coordination/InfractionCollector.sol
+++ b/contracts/contracts/coordination/InfractionCollector.sol
@@ -12,8 +12,13 @@ contract InfractionCollector is OwnableUpgradeable {
     // Reference to the TACoChildApplication contract
     ITACoChildApplication public tacoChildApplication;
 
+    // infraction types
+    enum InfractionType {
+        MISSING_TRANSCRIPT
+    }
+
     // Mapping to keep track of reported infractions
-    mapping(uint32 => mapping(address => bool)) public infractions;
+    mapping(uint32 => mapping(address => mapping(InfractionType => bool))) public infractions;
 
     function initialize(
         Coordinator _coordinator,
@@ -36,7 +41,7 @@ contract InfractionCollector is OwnableUpgradeable {
 
         for (uint256 i = 0; i < stakingProviders.length; i++) {
             // Check if the infraction has already been reported
-            require(!infractions[ritualId][stakingProviders[i]], "Infraction already reported");
+            require(!infractions[ritualId][stakingProviders[i]][InfractionType.MISSING_TRANSCRIPT], "Infraction already reported");
             Coordinator.Participant memory participant = coordinator.getParticipantFromProvider(
                 ritualId,
                 stakingProviders[i]
@@ -44,7 +49,7 @@ contract InfractionCollector is OwnableUpgradeable {
             if (participant.transcript.length == 0) {  // Transcript TX wasn't posted
                 // Penalize the staking provider
                 tacoChildApplication.penalize(stakingProviders[i]);
-                infractions[ritualId][stakingProviders[i]] = true;
+                infractions[ritualId][stakingProviders[i]][InfractionType.MISSING_TRANSCRIPT] = true;
             }
         }
     }
diff --git a/tests/test_infraction.py b/tests/test_infraction.py
index fe4de5e75..c9a166309 100644
--- a/tests/test_infraction.py
+++ b/tests/test_infraction.py
@@ -3,10 +3,6 @@
 
 import ape
 import pytest
-from eth_account import Account
-from eth_account.messages import encode_defunct
-from hexbytes import HexBytes
-from web3 import Web3
 
 TIMEOUT = 1000
 MAX_DKG_SIZE = 31
@@ -156,7 +152,7 @@ def test_report_infractions(erc20, nodes, initiator, global_allow_list, infracti
     chain.pending_timestamp += TIMEOUT * 2
     infraction_collector.reportMissingTranscript(0, nodes, sender=initiator)
     for node in nodes:
-        assert infraction_collector.infractions(0, node) == True
+        assert infraction_collector.infractions(0, node, 0) == True
 
 def test_cant_report_infractions_twice(erc20, nodes, initiator, global_allow_list, infraction_collector, coordinator, chain):
     for node in nodes: