Merge pull request #25 from flifloo/warn

Warn

administrator/utils.py

@@ -0,0 +1,25 @@
+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()
+
+
+def seconds_to_time_string(seconds: float) -> str:
+    days, seconds = divmod(seconds, 86400)
+    hours, seconds = divmod(seconds, 3600)
+    minutes, seconds = divmod(seconds, 60)
+    return f"""{f"{days}d {hours}h {minutes}m {seconds}s"
+    if days > 0 else f"{hours}h {minutes}m {seconds}s"
+    if hours > 0 else f"{minutes}m {seconds}s"
+    if minutes > 0 else f"{seconds}s"}"""

db/Warn.py

@@ -0,0 +1,22 @@
+from datetime import datetime
+
+from discord import Embed
+
+from db import Base
+from sqlalchemy import Column, Integer, BigInteger, Text, DateTime
+
+
+class Warn(Base):
+    __tablename__ = "warns"
+    id = Column(Integer, primary_key=True)
+    user = Column(BigInteger, nullable=False)
+    author = Column(BigInteger, nullable=False)
+    guild = Column(BigInteger, nullable=False)
+    description = Column(Text, nullable=False)
+    date = Column(DateTime, nullable=False, default=datetime.now())
+
+    def __init__(self, user: int, author: int, guild: int, description: str):
+        self.user = user
+        self.author = author
+        self.guild = guild
+        self.description = description

db/WarnAction.py

@@ -0,0 +1,19 @@
+from datetime import timedelta
+
+from db import Base
+from sqlalchemy import Column, Integer, BigInteger, String
+
+
+class WarnAction(Base):
+    __tablename__ = "warn_actions"
+    id = Column(Integer, primary_key=True)
+    guild = Column(BigInteger, nullable=False)
+    count = Column(Integer, nullable=False, unique=True)
+    action = Column(String, nullable=False)
+    duration = Column(BigInteger)
+
+    def __init__(self, guild: int, count: int, action: str, duration: float = None):
+        self.guild = guild
+        self.count = count
+        self.action = action
+        self.duration = duration

db/__init__.py

@@ -10,4 +10,6 @@ from db.Greetings import Greetings
 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)

extensions/__init__.py

@@ -8,3 +8,4 @@ bot.load_extension("extensions.reminder")
 bot.load_extension("extensions.greetings")
 bot.load_extension("extensions.presentation")
 bot.load_extension("extensions.rorec")
+bot.load_extension("extensions.warn")

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, seconds_to_time_string
 
 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
@@ -56,12 +45,7 @@ class Reminders(commands.Cog, name="Reminder"):
         s.commit()
         s.close()
 
-        hours, seconds = divmod(time.seconds, 3600)
-        minutes, seconds = divmod(seconds, 60)
-        await ctx.send(f"""Remind you in {f"{time.days}d {hours}h {minutes}m {seconds}s"
-        if time.days > 0 else f"{hours}h {minutes}m {seconds}s"
-        if hours > 0 else f"{minutes}m {seconds}s"
-        if minutes > 0 else f"{seconds}s"} !""")
+        await ctx.send(f"""Remind you in {seconds_to_time_string(time.total_seconds())} !""")
 
     @reminder.group("list", pass_context=True)
     async def reminder_list(self, ctx: commands.Context):

extensions/warn.py

@@ -0,0 +1,198 @@
+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, seconds_to_time_string
+
+extension_name = "warn"
+logger = logger.getChild(extension_name)
+
+
+class Warn(commands.Cog):
+    def __init__(self, bot: commands.Bot):
+        self.bot = bot
+
+    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}
+        if user not in users:
+            raise BadArgument()
+        return users[user]
+
+    @commands.group("warn", pass_context=True)
+    @commands.guild_only()
+    #@commands.has_permissions(manage_roles=True, kick_members=True, ban_members=True, mute_members=True)
+    async def warn(self, ctx: commands.Context):
+        if ctx.invoked_subcommand is None:
+            await ctx.invoke(self.warn_help)
+
+    @warn.group("help", pass_context=True)
+    async def warn_help(self, ctx: commands.Context):
+        embed = Embed(title="Warn help")
+        embed.add_field(name="add <user> <description>", value="Send a warn to a user", inline=False)
+        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#discriminator|actions]", value="List warn of the guild or a specified user\n"
+                                                                        "If you specify `actions` instead of a user, "
+                                                                        "all the actions of the guild will be listed",
+                        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 ban1D` to ban someone for one day after 3 "
+                                                              "warns", inline=False)
+        await ctx.send(embed=embed)
+
+    @warn.group("add", pass_context=True)
+    async def warn_add(self, ctx: commands.Context, user: str, description: str):
+        target = self.get_target(ctx, user)
+
+        s = db.Session()
+        s.add(db.Warn(target.id, ctx.author.id, ctx.guild.id, description))
+        s.commit()
+        s.close()
+
+        try:
+            embed = Embed(title="You get warned !", description="A moderator send you a warn", color=0xff0000)
+            embed.add_field(name="Description:", value=description)
+            await target.send(embed=embed)
+        except Forbidden:
+            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):
+        target = self.get_target(ctx, user)
+        s = db.Session()
+        ws = s.query(db.Warn).filter(db.Warn.guild == ctx.guild.id, db.Warn.user == target.id).all()
+        if number <= 0 or number > len(ws):
+            raise BadArgument()
+        s.delete(ws[number-1])
+        s.commit()
+        s.close()
+        await ctx.message.add_reaction("\U0001f44d")
+
+    @warn.group("purge", pass_context=True)
+    async def warn_purge(self, ctx: commands.Context, user: str):
+        target = self.get_target(ctx, user)
+        s = db.Session()
+        for w in s.query(db.Warn).filter(db.Warn.guild == ctx.guild.id, db.Warn.user == target.id).all():
+            s.delete(w)
+        s.commit()
+        s.close()
+        await ctx.message.add_reaction("\U0001f44d")
+
+    @warn.group("list", pass_context=True)
+    async def warn_list(self, ctx: commands.Context, user: str = None):
+        s = db.Session()
+        embed = Embed(title="Warn list")
+        ws = {}
+
+        if user == "actions":
+            embed.title = "Actions list"
+            for a in s.query(db.WarnAction).filter(db.WarnAction.guild == ctx.guild.id).order_by(db.WarnAction.count)\
+                    .all():
+                action = f"{a.action} for {seconds_to_time_string(a.duration)}" if a.duration else a.action
+                embed.add_field(name=f"{a.count} warn(s)", value=action, inline=False)
+        elif user:
+            target = self.get_target(ctx, user)
+            ws[target.id] = s.query(db.Warn).filter(db.Warn.guild == ctx.guild.id, db.Warn.user == target.id).all()
+        else:
+            for w in s.query(db.Warn).filter(db.Warn.guild == ctx.guild.id).all():
+                if w.user not in ws:
+                    ws[w.user] = []
+                ws[w.user].append(w)
+        s.close()
+
+        for u in ws:
+            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", "")).total_seconds()
+                action = "mute"
+            elif action.startswith("ban"):
+                if action[3:]:
+                    time = time_pars(action.replace("ban", "")).total_seconds()
+                    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()
+
+
+def setup(bot):
+    logger.info(f"Loading...")
+    try:
+        bot.add_cog(Warn(bot))
+    except Exception as e:
+        logger.error(f"Error loading: {e}")
+    else:
+        logger.info(f"Load successful")
+
+
+def teardown(bot):
+    logger.info(f"Unloading...")
+    try:
+        bot.remove_cog("Warn")
+    except Exception as e:
+        logger.error(f"Error unloading: {e}")
+    else:
+        logger.info(f"Unload successful")