From 75b516f986a1db467af8f0f7e601633eef85d6ce Mon Sep 17 00:00:00 2001
From: flifloo <flifloo@gmail.com>
Date: Sun, 2 Aug 2020 19:41:38 +0200
Subject: [PATCH] Add warn action command

---
 administrator/utils.py | 15 +++++++++++++
 db/WarnAction.py       | 20 +++++++++++++++++
 db/__init__.py         |  1 +
 extensions/reminder.py | 13 +----------
 extensions/warn.py     | 49 +++++++++++++++++++++++++++++++++++++-----
 5 files changed, 81 insertions(+), 17 deletions(-)
 create mode 100644 administrator/utils.py
 create mode 100644 db/WarnAction.py

diff --git a/administrator/utils.py b/administrator/utils.py
new file mode 100644
index 0000000..32e609f
--- /dev/null
+++ b/administrator/utils.py
@@ -0,0 +1,15 @@
+import re
+from datetime import timedelta
+
+from discord.ext.commands import BadArgument
+
+
+def time_pars(s: str) -> timedelta:
+    match = re.fullmatch(r"(?:([0-9]+)W)*(?:([0-9]+)D)*(?:([0-9]+)H)*(?:([0-9]+)M)*(?:([0-9]+)S)*",
+                         s.upper().replace(" ", "").strip())
+    if match:
+        w, d, h, m, s = match.groups()
+        if any([w, d, h, m, s]):
+            w, d, h, m, s = [i if i else 0 for i in [w, d, h, m, s]]
+            return timedelta(weeks=int(w), days=int(d), hours=int(h), minutes=int(m), seconds=int(s))
+    raise BadArgument()
diff --git a/db/WarnAction.py b/db/WarnAction.py
new file mode 100644
index 0000000..f923410
--- /dev/null
+++ b/db/WarnAction.py
@@ -0,0 +1,20 @@
+from datetime import timedelta
+
+from db import Base
+from sqlalchemy import Column, Integer, BigInteger, Float, String
+
+
+class WarnAction(Base):
+    __tablename__ = "warn_actions"
+    id = Column(Integer, primary_key=True)
+    guild = Column(BigInteger, nullable=False)
+    count = Column(Float, nullable=False, unique=True)
+    action = Column(String, nullable=False)
+    duration = Column(BigInteger)
+
+    def __init__(self, guild: int, count: int, action: str, duration: timedelta = None):
+        self.guild = guild
+        self.count = count
+        self.action = action
+        if duration:
+            self.duration = duration.total_seconds()
diff --git a/db/__init__.py b/db/__init__.py
index fd2e1db..904b1f0 100644
--- a/db/__init__.py
+++ b/db/__init__.py
@@ -11,4 +11,5 @@ from db.Presentation import Presentation
 from db.RoRec import RoRec
 from db.Polls import Polls
 from db.Warn import Warn
+from db.WarnAction import WarnAction
 Base.metadata.create_all(engine)
diff --git a/extensions/reminder.py b/extensions/reminder.py
index aef4310..d2b8913 100644
--- a/extensions/reminder.py
+++ b/extensions/reminder.py
@@ -8,23 +8,12 @@ from discord.ext import tasks
 
 from administrator.logger import logger
 from administrator import db
-
+from administrator.utils import time_pars
 
 extension_name = "reminders"
 logger = logger.getChild(extension_name)
 
 
-def time_pars(s: str) -> timedelta:
-    match = re.fullmatch(r"(?:([0-9]+)W)*(?:([0-9]+)D)*(?:([0-9]+)H)*(?:([0-9]+)M)*(?:([0-9]+)S)*",
-                         s.upper().replace(" ", "").strip())
-    if match:
-        w, d, h, m, s = match.groups()
-        if any([w, d, h, m, s]):
-            w, d, h, m, s = [i if i else 0 for i in [w, d, h, m, s]]
-            return timedelta(weeks=int(w), days=int(d), hours=int(h), minutes=int(m), seconds=int(s))
-    raise BadArgument()
-
-
 class Reminders(commands.Cog, name="Reminder"):
     def __init__(self, bot: commands.Bot):
         self.bot = bot
diff --git a/extensions/warn.py b/extensions/warn.py
index b604d7e..c9b1296 100644
--- a/extensions/warn.py
+++ b/extensions/warn.py
@@ -1,17 +1,14 @@
-import re
-
 from discord import Embed, Forbidden, Member, Guild
 from discord.ext import commands
 from discord.ext.commands import BadArgument
 
 from administrator import db
 from administrator.logger import logger
+from administrator.utils import time_pars
 
 extension_name = "warn"
 logger = logger.getChild(extension_name)
 
-channel_id_re = re.compile(r"^<#([0-9]+)>$")
-
 
 class Warn(commands.Cog):
     def __init__(self, bot: commands.Bot):
@@ -41,6 +38,13 @@ class Warn(commands.Cog):
         embed.add_field(name="remove <user> <number>", value="Remove a number of warn to a user", inline=False)
         embed.add_field(name="purge <user>", value="Remove all warn of a user", inline=False)
         embed.add_field(name="list [user]", value="List warn of the guild or a specified user", inline=False)
+        embed.add_field(name="action <count> <action>", value="Set an action for a count of warn\n"
+                                                              "Actions: `mute<time>`, `kick`, `ban[time]`, `nothing`\n"
+                                                              "Time: `?D?H?M?S`\n"
+                                                              "Example: `action 1 mute1H` to mute someone for one hour "
+                                                              "after only one war\n"
+                                                              "or `action 3 ban3D` to ban someone for one day after 3 "
+                                                              "warns", inline=False)
         await ctx.send(embed=embed)
 
     @warn.group("add", pass_context=True)
@@ -100,16 +104,51 @@ class Warn(commands.Cog):
         s.close()
 
         for u in ws:
-            warns = [f"{w.date.strftime('%d/%m/%Y %H:%M')} - {w.description}" for w in ws[u]]
+            warns = [f"{self.bot.get_user(w.author).mention} - {w.date.strftime('%d/%m/%Y %H:%M')}```{w.description}```"
+                     for w in ws[u]]
             embed.add_field(name=self.bot.get_user(u), value="\n".join(warns), inline=False)
 
         await ctx.send(embed=embed)
 
+    @warn.group("action", pass_context=True)
+    async def warn_action(self, ctx: commands.Context, count: int, action: str):
+        if count <= 0 or\
+                (action not in ["kick", "nothing"] and not action.startswith("mute") and not action.startswith("ban")):
+            raise BadArgument()
+
+        s = db.Session()
+        a = s.query(db.WarnAction).filter(db.WarnAction.guild == ctx.guild.id, db.WarnAction.count == count).first()
+
+        if action == "nothing":
+            if a:
+                s.delete(a)
+            else:
+                raise BadArgument()
+        else:
+            time = None
+            if action.startswith("mute"):
+                time = time_pars(action.replace("mute", ""))
+                action = "mute"
+            elif action.startswith("ban"):
+                time = time_pars(action.replace("ban", ""))
+                action = "ban"
+            if a:
+                a.action = action
+                a.duration = time
+            else:
+                s.add(db.WarnAction(ctx.guild.id, count, action, time))
+
+        s.commit()
+        s.close()
+        await ctx.message.add_reaction("\U0001f44d")
+
     @commands.Cog.listener()
     async def on_guild_remove(self, guild: Guild):
         s = db.Session()
         for w in s.query(db.Warn).filter(db.Warn.guild == guild.id).all():
             s.delete(w)
+        for a in s.query(db.WarnAction).filter(db.WarnAction.guild == guild.id).all():
+            s.delete(a)
         s.commit()
         s.close()