From 26dda9142d305c58cf41f3a97c856defb1514d88 Mon Sep 17 00:00:00 2001
From: flifloo <flifloo@gmail.com>
Date: Sun, 2 Aug 2020 20:13:18 +0200
Subject: [PATCH] Check if action need after each warn add

---
 db/WarnAction.py   |  5 ++---
 extensions/warn.py | 22 +++++++++++++++++++---
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/db/WarnAction.py b/db/WarnAction.py
index f923410..634e29d 100644
--- a/db/WarnAction.py
+++ b/db/WarnAction.py
@@ -12,9 +12,8 @@ class WarnAction(Base):
     action = Column(String, nullable=False)
     duration = Column(BigInteger)
 
-    def __init__(self, guild: int, count: int, action: str, duration: timedelta = None):
+    def __init__(self, guild: int, count: int, action: str, duration: float = None):
         self.guild = guild
         self.count = count
         self.action = action
-        if duration:
-            self.duration = duration.total_seconds()
+        self.duration = duration
diff --git a/extensions/warn.py b/extensions/warn.py
index c9b1296..e481856 100644
--- a/extensions/warn.py
+++ b/extensions/warn.py
@@ -17,6 +17,20 @@ class Warn(commands.Cog):
     def description(self):
         return "Send warning to user and make custom action after a number of warn"
 
+    @staticmethod
+    async def check_warn(ctx: commands.Context, target: Member):
+        s = db.Session()
+        c = s.query(db.Warn).filter(db.Warn.guild == ctx.guild.id, db.Warn.user == target.id).count()
+        a = s.query(db.WarnAction).filter(db.WarnAction.guild == ctx.guild.id, db.WarnAction.count == c).first()
+        if a:
+            reason = f"Action after {c} warns"
+            if a.action == "kick":
+                await target.kick(reason=reason)
+            elif a.action == "ban":
+                await target.ban(reason=reason)
+            elif a.action == "mute":
+                pass  # Integration with upcoming ban & mute extension
+
     @staticmethod
     def get_target(ctx: commands.Context, user: str) -> Member:
         users = {str(m): m for m in ctx.guild.members}
@@ -64,6 +78,7 @@ class Warn(commands.Cog):
             await ctx.send("Fail to send warn notification to the user, DM close :warning:")
         else:
             await ctx.message.add_reaction("\U0001f44d")
+        await self.check_warn(ctx, target)
 
     @warn.group("remove", pass_context=True)
     async def warn_remove(self, ctx: commands.Context, user: str, number: int):
@@ -127,11 +142,12 @@ class Warn(commands.Cog):
         else:
             time = None
             if action.startswith("mute"):
-                time = time_pars(action.replace("mute", ""))
+                time = time_pars(action.replace("mute", "")).total_seconds()
                 action = "mute"
             elif action.startswith("ban"):
-                time = time_pars(action.replace("ban", ""))
-                action = "ban"
+                if action[3:]:
+                    time = time_pars(action.replace("ban", "")).total_seconds()
+                    action = "ban"
             if a:
                 a.action = action
                 a.duration = time