Fix equip packets, again

src/main/java/me/libraryaddict/disguise/utilities/packets/LibsPackets.java

@@ -3,11 +3,18 @@ package me.libraryaddict.disguise.utilities.packets;
 import com.comphenix.protocol.PacketType;
 import com.comphenix.protocol.ProtocolLibrary;
 import com.comphenix.protocol.events.PacketContainer;
+import com.comphenix.protocol.reflect.StructureModifier;
+import lombok.Getter;
+import lombok.Setter;
 import me.libraryaddict.disguise.LibsDisguises;
 import me.libraryaddict.disguise.disguisetypes.Disguise;
-import me.libraryaddict.disguise.utilities.DisguiseUtilities;
+import me.libraryaddict.disguise.utilities.reflection.ReflectionManager;
 import org.bukkit.Bukkit;
+import org.bukkit.Material;
+import org.bukkit.entity.LivingEntity;
 import org.bukkit.entity.Player;
+import org.bukkit.inventory.EquipmentSlot;
+import org.bukkit.inventory.ItemStack;
 
 import java.lang.reflect.InvocationTargetException;
 import java.util.*;
@@ -22,6 +29,9 @@ public class LibsPackets {
     private Disguise disguise;
     private boolean doNothing;
     private int removeMetaAt = -1;
+    @Getter
+    @Setter
+    private boolean sendArmor;
 
     public LibsPackets(Disguise disguise) {
         this.disguise = disguise;
@@ -83,8 +93,33 @@ public class LibsPackets {
             final boolean isRemoveCancel = isSpawnPacket && entry.getKey() >= removeMetaAt && removeMetaAt >= 0;
 
             Bukkit.getScheduler().scheduleSyncDelayedTask(LibsDisguises.getInstance(), () -> {
+                if (!disguise.isDisguiseInUse()) {
+                    if (isRemoveCancel) {
+                        PacketsManager.getPacketsHandler().removeCancel(disguise, observer);
+                    }
+
+                    return;
+                }
+
                 if (isRemoveCancel) {
                     PacketsManager.getPacketsHandler().removeCancel(disguise, observer);
+
+                    if (isSendArmor()) {
+                        for (EquipmentSlot slot : EquipmentSlot.values()) {
+                            PacketContainer packet = createPacket(slot);
+
+                            if (packet == null) {
+                                continue;
+                            }
+
+                            try {
+                                ProtocolLibrary.getProtocolManager().sendServerPacket(observer, packet, false);
+                            }
+                            catch (InvocationTargetException e) {
+                                e.printStackTrace();
+                            }
+                        }
+                    }
                 }
 
                 try {
@@ -98,4 +133,31 @@ public class LibsPackets {
             }, entry.getKey());
         }
     }
+
+    private PacketContainer createPacket(EquipmentSlot slot) {
+        // Get what the disguise wants to show for its armor
+        ItemStack itemToSend = disguise.getWatcher().getItemStack(slot);
+
+        // If the disguise armor isn't visible
+        if (itemToSend == null) {
+            itemToSend = ReflectionManager.getEquipment(slot, disguise.getEntity());
+
+            // If natural armor isn't sent either
+            if (itemToSend == null || itemToSend.getType() == Material.AIR) {
+                return null;
+            }
+        } else if (itemToSend.getType() == Material.AIR) {
+            return null;
+        }
+
+        PacketContainer packet = new PacketContainer(PacketType.Play.Server.ENTITY_EQUIPMENT);
+
+        StructureModifier<Object> mods = packet.getModifier();
+
+        mods.write(0, disguise.getEntity().getEntityId());
+        mods.write(1, ReflectionManager.createEnumItemSlot(slot));
+        mods.write(2, ReflectionManager.getNmsItem(itemToSend));
+
+        return packet;
+    }
 }