fix: restore abilities when toggling

1 files changed, 18 insertions, 4 deletions

diff --git a/src/main/java/co/fantabos/bscv/modules/VanillaFlight.java b/src/main/java/co/fantabos/bscv/modules/VanillaFlight.java
index f8e2312..767bfff 100644
--- a/src/main/java/co/fantabos/bscv/modules/VanillaFlight.java
+++ b/src/main/java/co/fantabos/bscv/modules/VanillaFlight.java
@@ -4,7 +4,6 @@ import com.mojang.brigadier.CommandDispatcher;
 
 import co.fantabos.bscv.BoSCoVicino;
 import co.fantabos.bscv.Module;
-import net.minecraft.client.entity.player.ClientPlayerEntity;
 import net.minecraft.command.CommandSource;
 import net.minecraftforge.common.ForgeConfigSpec;
 import net.minecraftforge.event.TickEvent;
@@ -16,11 +15,26 @@ public class VanillaFlight extends Module {
 		super("VanillaFlight", Group.CORE, builder, dispatcher);
 	}
 
+	private boolean couldFlyBefore = false;
+
 	@SubscribeEvent
 	public void onTick(TickEvent.ClientTickEvent event) {
-		ClientPlayerEntity player = BoSCoVicino.minecraft.player;
-		if (player != null) {
-			player.abilities.mayfly = true;
+		if (BoSCoVicino.minecraft.player != null) {
+			BoSCoVicino.minecraft.player.abilities.mayfly = true;
+		}
+	}
+
+	@Override
+	protected void onEnabled() {
+		if (BoSCoVicino.minecraft.player != null) {
+			this.couldFlyBefore = BoSCoVicino.minecraft.player.abilities.mayfly;
+		}
+	}
+
+	@Override
+	protected void onDisabled() {
+		if (BoSCoVicino.minecraft.player != null) {
+			BoSCoVicino.minecraft.player.abilities.mayfly = this.couldFlyBefore;
 		}
 	}
 }