1 files changed, 21 insertions, 1 deletions
diff --git a/src/main/java/world/anhgelus/lifesteal/LifeSteal.java b/src/main/java/world/anhgelus/lifesteal/LifeSteal.java
index b237e6c..062278d 100644
--- a/src/main/java/world/anhgelus/lifesteal/LifeSteal.java
+++ b/src/main/java/world/anhgelus/lifesteal/LifeSteal.java
@@ -19,6 +19,7 @@ import net.minecraft.server.command.CommandManager;
 import net.minecraft.server.command.ServerCommandSource;
 import net.minecraft.server.network.ServerPlayNetworkHandler;
 import net.minecraft.server.network.ServerPlayerEntity;
+import net.minecraft.stat.Stats;
 import net.minecraft.util.ActionResult;
 import net.minecraft.world.GameRules;
 import org.slf4j.Logger;
@@ -31,13 +32,21 @@ public class LifeSteal implements ModInitializer {
     public static final Logger LOGGER = LoggerFactory.getLogger(MOD_ID);
     public static float MIN_HEALTH = 0;
     public static float MAX_HEALTH = 40;
-    public static final GameRules.Key<GameRules.IntRule> MAX_HEARTS = GameRuleRegistry.register(
+    public static final GameRules.Key<GameRules.IntRule> RULE_MAX_HEARTS = GameRuleRegistry.register(
             MOD_ID + ":maxHearts",
             GameRules.Category.MISC,
             GameRuleFactory.createIntRule((int) MAX_HEALTH / 2, (server, rule) -> {
                 MAX_HEALTH = rule.get() * 2;
             })
     );
+    public static int GRACE_PERIOD = 30;
+    public static final GameRules.Key<GameRules.IntRule> RULE_GRACE_PERIOD = GameRuleRegistry.register(
+            MOD_ID + ":gracePeriod",
+            GameRules.Category.MISC,
+            GameRuleFactory.createIntRule(GRACE_PERIOD, (server, rule) -> {
+                GRACE_PERIOD = rule.get();
+            })
+    );
 
     @Override
     public void onInitialize() {
@@ -67,6 +76,17 @@ public class LifeSteal implements ModInitializer {
             if (damageSource.getAttacker() instanceof ServerPlayerEntity killer) getLifeStealer(killer).kill();
         });
 
+        final var playTimeStat = Stats.CUSTOM.getOrCreateStat(Stats.PLAY_TIME);
+
+        ServerLivingEntityEvents.ALLOW_DAMAGE.register((entity, source, amount) -> {
+            if (!(entity instanceof ServerPlayerEntity player)) return true;
+            if (!(source.getAttacker() instanceof ServerPlayerEntity attacker)) return true;
+            final var playTime = player.getStatHandler().getStat(playTimeStat);
+            final var attackerPlayTime = attacker.getStatHandler().getStat(playTimeStat);
+            final var grace = GRACE_PERIOD * 60 * 20;
+            return playTime > grace && attackerPlayTime > grace;
+        });
+
         ServerPlayerEvents.AFTER_RESPAWN.register((ServerPlayerEntity oldPlayer, ServerPlayerEntity newPlayer, boolean alive) -> {
             getLifeStealer(oldPlayer).respawn(newPlayer);
         });