diff --git a/OpenRA.Mods.Common/UtilityCommands/UpgradeRules.cs b/OpenRA.Mods.Common/UtilityCommands/UpgradeRules.cs
index d8cc9842cc..b40f38561e 100644
--- a/OpenRA.Mods.Common/UtilityCommands/UpgradeRules.cs
+++ b/OpenRA.Mods.Common/UtilityCommands/UpgradeRules.cs
@@ -637,6 +637,13 @@ namespace OpenRA.Mods.Common.UtilityCommands
 						if (!node.Value.Nodes.Any(n => n.Key == "Condition"))
 							node.Value.Nodes.Add(new MiniYamlNode("Condition", "terrain"));
 					}
+
+					if (node.Key == "AttackSwallow")
+					{
+						ConvertUpgradesToCondition(parent, node, "AttackingUpgrades", "AttackingCondition");
+						if (!node.Value.Nodes.Any(n => n.Key == "AttackingCondition"))
+							node.Value.Nodes.Add(new MiniYamlNode("AttackingCondition", "attacking"));
+					}
 				}
 
 				UpgradeActorRules(modData, engineVersion, ref node.Value.Nodes, node, depth + 1);
diff --git a/OpenRA.Mods.D2k/Activities/SwallowActor.cs b/OpenRA.Mods.D2k/Activities/SwallowActor.cs
index a8344700c6..a5f220c772 100644
--- a/OpenRA.Mods.D2k/Activities/SwallowActor.cs
+++ b/OpenRA.Mods.D2k/Activities/SwallowActor.cs
@@ -37,6 +37,7 @@ namespace OpenRA.Mods.D2k.Activities
 		int countdown;
 		CPos burrowLocation;
 		AttackState stance;
+		int attackingToken = UpgradeManager.InvalidConditionToken;
 
 		public SwallowActor(Actor self, Target target, WeaponInfo weapon)
 		{
@@ -45,7 +46,7 @@ namespace OpenRA.Mods.D2k.Activities
 			sandworm = self.Trait<Sandworm>();
 			positionable = self.Trait<Mobile>();
 			swallow = self.Trait<AttackSwallow>();
-			manager = self.Trait<UpgradeManager>();
+			manager = self.TraitOrDefault<UpgradeManager>();
 			radarPings = self.World.WorldActor.TraitOrDefault<RadarPings>();
 		}
 
@@ -105,10 +106,12 @@ namespace OpenRA.Mods.D2k.Activities
 			switch (stance)
 			{
 				case AttackState.Uninitialized:
-					GrantUpgrades(self);
 					stance = AttackState.Burrowed;
 					countdown = swallow.Info.AttackDelay;
 					burrowLocation = self.Location;
+					if (manager != null && attackingToken == UpgradeManager.InvalidConditionToken &&
+							!string.IsNullOrEmpty(swallow.Info.AttackingCondition))
+						attackingToken = manager.GrantCondition(self, swallow.Info.AttackingCondition);
 					break;
 				case AttackState.Burrowed:
 					if (--countdown > 0)
@@ -119,14 +122,14 @@ namespace OpenRA.Mods.D2k.Activities
 					// The target has moved too far away
 					if ((burrowLocation - targetLocation).Length > NearEnough)
 					{
-						RevokeUpgrades(self);
+						RevokeCondition(self);
 						return NextActivity;
 					}
 
 					// The target reached solid ground
 					if (!positionable.CanEnterCell(targetLocation, null, false))
 					{
-						RevokeUpgrades(self);
+						RevokeCondition(self);
 						return NextActivity;
 					}
 
@@ -135,7 +138,7 @@ namespace OpenRA.Mods.D2k.Activities
 
 					if (!targets.Any())
 					{
-						RevokeUpgrades(self);
+						RevokeCondition(self);
 						return NextActivity;
 					}
 
@@ -158,23 +161,17 @@ namespace OpenRA.Mods.D2k.Activities
 						self.World.AddFrameEndTask(w => self.Dispose());
 					}
 
-					RevokeUpgrades(self);
+					RevokeCondition(self);
 					return NextActivity;
 			}
 
 			return this;
 		}
 
-		void GrantUpgrades(Actor self)
+		void RevokeCondition(Actor self)
 		{
-			foreach (var up in swallow.Info.AttackingUpgrades)
-				manager.GrantUpgrade(self, up, this);
-		}
-
-		void RevokeUpgrades(Actor self)
-		{
-			foreach (var up in swallow.Info.AttackingUpgrades)
-				manager.RevokeUpgrade(self, up, this);
+			if (attackingToken != UpgradeManager.InvalidConditionToken)
+				attackingToken = manager.RevokeCondition(self, attackingToken);
 		}
 	}
 }
diff --git a/OpenRA.Mods.D2k/Traits/AttackSwallow.cs b/OpenRA.Mods.D2k/Traits/AttackSwallow.cs
index 3516680a71..19b8fc53f7 100644
--- a/OpenRA.Mods.D2k/Traits/AttackSwallow.cs
+++ b/OpenRA.Mods.D2k/Traits/AttackSwallow.cs
@@ -27,8 +27,8 @@ namespace OpenRA.Mods.D2k.Traits
 		public readonly int AttackDelay = 30;
 
 		[UpgradeGrantedReference]
-		[Desc("The upgrades to grant while attacking.")]
-		public readonly string[] AttackingUpgrades = { "attacking" };
+		[Desc("The condition to grant to self while attacking.")]
+		public readonly string AttackingCondition = null;
 
 		public readonly string WormAttackSound = "WORM.WAV";
 
diff --git a/mods/d2k/rules/arrakis.yaml b/mods/d2k/rules/arrakis.yaml
index 1fec4f6d62..680f14e3cd 100644
--- a/mods/d2k/rules/arrakis.yaml
+++ b/mods/d2k/rules/arrakis.yaml
@@ -87,6 +87,7 @@ sandworm:
 	AttackSwallow:
 		AttackRequiresEnteringCell: true
 		IgnoresVisibility: true
+		AttackingCondition: attacking
 	Armament:
 		Weapon: WormJaw
 	Sandworm: