From 084b482cfe9ae8f54a7015ff3ee14c71cfe6c49f Mon Sep 17 00:00:00 2001
From: Rane <60792108+Elijahrane@users.noreply.github.com>
Date: Thu, 3 Nov 2022 21:37:17 -0400
Subject: [PATCH] Entity whitelist/blacklist for traits (#12348)

---
 Content.Server/Traits/TraitSystem.cs           |  7 +++++++
 Content.Shared/Traits/TraitPrototype.cs        | 13 +++++++++++++
 Resources/Prototypes/Traits/disabilities.yml   |  3 +++
 Resources/Prototypes/Traits/inconveniences.yml |  3 +++
 4 files changed, 26 insertions(+)

diff --git a/Content.Server/Traits/TraitSystem.cs b/Content.Server/Traits/TraitSystem.cs
index c0c264ed17..ceec4c9280 100644
--- a/Content.Server/Traits/TraitSystem.cs
+++ b/Content.Server/Traits/TraitSystem.cs
@@ -1,5 +1,6 @@
 using Content.Server.GameTicking;
 using Content.Shared.Traits;
+using Content.Shared.Whitelist;
 using Robust.Shared.Prototypes;
 using Robust.Shared.Serialization.Manager;
 
@@ -28,6 +29,12 @@ public sealed class TraitSystem : EntitySystem
                 return;
             }
 
+            if (traitPrototype.Whitelist != null && !traitPrototype.Whitelist.IsValid(args.Mob))
+                continue;
+
+            if (traitPrototype.Blacklist != null && traitPrototype.Blacklist.IsValid(args.Mob))
+                continue;
+
             // Add all components required by the prototype
             foreach (var entry in traitPrototype.Components.Values)
             {
diff --git a/Content.Shared/Traits/TraitPrototype.cs b/Content.Shared/Traits/TraitPrototype.cs
index 94f0b39c0a..1b3a73ce6b 100644
--- a/Content.Shared/Traits/TraitPrototype.cs
+++ b/Content.Shared/Traits/TraitPrototype.cs
@@ -1,3 +1,4 @@
+using Content.Shared.Whitelist;
 using Robust.Shared.Prototypes;
 using static Robust.Shared.Prototypes.EntityPrototype; // don't worry about it
 
@@ -25,6 +26,18 @@ namespace Content.Shared.Traits
         [DataField("description")]
         public string? Description { get; }
 
+        /// <summary>
+        ///     Don't apply this trait to entities this whitelist IS NOT valid for.
+        /// </summary>
+        [DataField("whitelist")]
+        public EntityWhitelist? Whitelist;
+
+        /// <summary>
+        ///     Don't apply this trait to entities this whitelist IS valid for. (hence, a blacklist)
+        /// </summary>
+        [DataField("blacklist")]
+        public EntityWhitelist? Blacklist;
+
         /// <summary>
         ///     The components that get added to the player, when they pick this trait.
         /// </summary>
diff --git a/Resources/Prototypes/Traits/disabilities.yml b/Resources/Prototypes/Traits/disabilities.yml
index e6dfd5a9b9..4746013445 100644
--- a/Resources/Prototypes/Traits/disabilities.yml
+++ b/Resources/Prototypes/Traits/disabilities.yml
@@ -2,6 +2,9 @@
   id: Blindness
   name: Blindness
   description: You lack vision
+  whitelist:
+    components:
+      - Blindable
   components:
     - type: PermanentBlindness
 
diff --git a/Resources/Prototypes/Traits/inconveniences.yml b/Resources/Prototypes/Traits/inconveniences.yml
index e8238c9bc1..94368b1519 100644
--- a/Resources/Prototypes/Traits/inconveniences.yml
+++ b/Resources/Prototypes/Traits/inconveniences.yml
@@ -2,6 +2,9 @@
   id: UncontrollableSneezing
   name: Runny nose
   description: You sneeze and cough uncontrollably
+  whitelist:
+    components:
+      - DiseaseCarrier
   components:
     - type: UncontrollableSnough
       snoughSound: