From 8ba8b7c0c68726972e398c33946aea081337515a Mon Sep 17 00:00:00 2001
From: Evil Factory <36804725+evilfactory@users.noreply.github.com>
Date: Mon, 11 Oct 2021 12:36:22 -0300
Subject: [PATCH] traitor hooks

---
 .../ServerSource/Traitors/TraitorMission.cs               | 8 ++++++++
 Barotrauma/BarotraumaShared/SharedSource/Lua/LuaSetup.cs  | 5 ++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/Barotrauma/BarotraumaServer/ServerSource/Traitors/TraitorMission.cs b/Barotrauma/BarotraumaServer/ServerSource/Traitors/TraitorMission.cs
index 1e16d660f..a6ed7337f 100644
--- a/Barotrauma/BarotraumaServer/ServerSource/Traitors/TraitorMission.cs
+++ b/Barotrauma/BarotraumaServer/ServerSource/Traitors/TraitorMission.cs
@@ -92,6 +92,13 @@ namespace Barotrauma
                 var traitorCandidates = new List<Tuple<Client, Character>>();
                 foreach (Client c in server.ConnectedClients)
                 {
+                    var result = new LuaResult(GameMain.Lua.hook.Call("traitor.findTraitorCandidate", new object[] { c, team }));
+					if (result.Bool())
+					{
+                        traitorCandidates.Add(Tuple.Create(c, c.Character));
+                        continue;
+                    }
+
                     if (c.Character == null || c.Character.IsDead || c.Character.Removed || !traitorRoleFilter(c.Character) ||
                         (team != CharacterTeamType.None && c.Character.TeamID != team))
                     {
@@ -227,6 +234,7 @@ namespace Barotrauma
                 foreach (var traitor in Traitors.Values)
                 {
                     traitor.Greet(server, CodeWords, CodeResponse, message => pendingMessages[traitor].Add(message));
+                    GameMain.Lua.hook.Call("traitor.traitorAssigned", new object[] { traitor });
                 }
                 pendingMessages.ForEach(traitor => traitor.Value.ForEach(message => traitor.Key.SendChatMessage(message, Identifier)));
                 pendingMessages.ForEach(traitor => traitor.Value.ForEach(message => traitor.Key.SendChatMessageBox(message, Identifier)));
diff --git a/Barotrauma/BarotraumaShared/SharedSource/Lua/LuaSetup.cs b/Barotrauma/BarotraumaShared/SharedSource/Lua/LuaSetup.cs
index 5253c72af..daa6589fc 100644
--- a/Barotrauma/BarotraumaShared/SharedSource/Lua/LuaSetup.cs
+++ b/Barotrauma/BarotraumaShared/SharedSource/Lua/LuaSetup.cs
@@ -392,6 +392,8 @@ namespace Barotrauma
 			UserData.RegisterType<Rand.RandSync>();
 			UserData.RegisterType<Skill>();
 			UserData.RegisterType<SkillPrefab>();
+			UserData.RegisterType<TraitorMissionPrefab>();
+			UserData.RegisterType<TraitorMissionResult>();
 
 			UserData.RegisterType<World>();
 			UserData.RegisterType<Fixture>();
@@ -412,7 +414,8 @@ namespace Barotrauma
 			AddCallMetaMember(UserData.RegisterType<SubmarineInfo>());
 
 #if SERVER
-
+			UserData.RegisterType<Traitor>();
+			UserData.RegisterType<Traitor.TraitorMission>();
 #elif CLIENT
 			UserData.RegisterType<LuaGUI>();
 			UserData.RegisterType<ChatBox>();