From e0fc6b0e3960cf4ecc5604992a2fef801973eb04 Mon Sep 17 00:00:00 2001
From: SmArtKar <44720187+SmArtKar@users.noreply.github.com>
Date: Sun, 15 May 2022 21:41:36 +0300
Subject: [PATCH] Fixes lightning bolt not giving tesla shock immunity when
 granted by an admin (#66973)

---
 code/modules/antagonists/wizard/equipment/spellbook.dm | 8 --------
 code/modules/spells/spell_types/aimed.dm               | 8 ++++++++
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/code/modules/antagonists/wizard/equipment/spellbook.dm b/code/modules/antagonists/wizard/equipment/spellbook.dm
index 62fa07028d6..98cb9f5451e 100644
--- a/code/modules/antagonists/wizard/equipment/spellbook.dm
+++ b/code/modules/antagonists/wizard/equipment/spellbook.dm
@@ -255,14 +255,6 @@
 	spell_type = /obj/effect/proc_holder/spell/aimed/lightningbolt
 	cost = 1
 
-/datum/spellbook_entry/lightningbolt/Buy(mob/living/carbon/human/user,obj/item/spellbook/book) //return TRUE on success
-	. = ..()
-	ADD_TRAIT(user, TRAIT_TESLA_SHOCKIMMUNE, "lightning_bolt_spell")
-
-/datum/spellbook_entry/lightningbolt/Refund(mob/living/carbon/human/user, obj/item/spellbook/book)
-	. = ..()
-	REMOVE_TRAIT(user, TRAIT_TESLA_SHOCKIMMUNE, "lightning_bolt_spell")
-
 /datum/spellbook_entry/infinite_guns
 	name = "Lesser Summon Guns"
 	desc = "Why reload when you have infinite guns? Summons an unending stream of bolt action rifles that deal little damage, but will knock targets down. Requires both hands free to use. Learning this spell makes you unable to learn Arcane Barrage."
diff --git a/code/modules/spells/spell_types/aimed.dm b/code/modules/spells/spell_types/aimed.dm
index 3e6e402675a..64ca91f0541 100644
--- a/code/modules/spells/spell_types/aimed.dm
+++ b/code/modules/spells/spell_types/aimed.dm
@@ -115,6 +115,14 @@
 	deactive_msg = "You let the energy flow out of your hands back into yourself..."
 	projectile_type = /obj/projectile/magic/aoe/lightning
 
+/obj/effect/proc_holder/spell/aimed/lightningbolt/on_gain(mob/living/user)
+	. = ..()
+	ADD_TRAIT(user, TRAIT_TESLA_SHOCKIMMUNE, "lightning_bolt_spell")
+
+/obj/effect/proc_holder/spell/aimed/lightningbolt/on_lose(mob/living/user)
+	. = ..()
+	REMOVE_TRAIT(user, TRAIT_TESLA_SHOCKIMMUNE, "lightning_bolt_spell")
+
 /obj/effect/proc_holder/spell/aimed/fireball
 	name = "Fireball"
 	desc = "This spell fires an explosive fireball at a target."