Merge pull request #6 from multitheftauto/tblock

Add duration (by default) to tblock command

Author: Lpsd

forward/converters.py

@@ -0,0 +1,56 @@
+import logging
+import re
+from typing import Union, Dict
+from datetime import timedelta
+
+from discord.ext.commands.converter import Converter
+from redbot.core import commands
+from redbot.core import i18n
+
+log = logging.getLogger("red.cogs.forward")
+
+# the following regex is slightly modified from Red
+# it's changed to be slightly more strict on matching with finditer
+# this is to prevent "empty" matches when parsing the full reason
+# This is also designed more to allow time interval at the beginning or the end of the mute
+# to account for those times when you think of adding time *after* already typing out the reason
+# https://github.com/Cog-Creators/Red-DiscordBot/blob/V3/develop/redbot/core/commands/converter.py#L55
+TIME_RE_STRING = r"|".join(
+    [
+        r"((?P<weeks>\d+?)\s?(weeks?|w))",
+        r"((?P<days>\d+?)\s?(days?|d))",
+        r"((?P<hours>\d+?)\s?(hours?|hrs|hr?))",
+        r"((?P<minutes>\d+?)\s?(minutes?|mins?|m(?!o)))",  # prevent matching "months"
+        r"((?P<seconds>\d+?)\s?(seconds?|secs?|s))",
+    ]
+)
+TIME_RE = re.compile(TIME_RE_STRING, re.I)
+TIME_SPLIT = re.compile(r"t(?:ime)?=")
+
+_ = i18n.Translator("Forward", __file__)
+
+def str_to_timedelta(
+    duration: str
+) -> Dict[str, Union[timedelta, str, None]]:
+    time_split = TIME_SPLIT.split(duration)
+    result: Dict[str, Union[timedelta, str, None]] = {}
+    if time_split:
+        maybe_time = time_split[-1]
+    else:
+        maybe_time = duration
+
+    time_data = {}
+    for time in TIME_RE.finditer(maybe_time):
+        duration = duration.replace(time[0], "")
+        for k, v in time.groupdict().items():
+            if v:
+                time_data[k] = int(v)
+    if time_data:
+        try:
+            result["duration"] = timedelta(**time_data)
+        except OverflowError:
+            raise commands.BadArgument(
+                _("The time provided is too long; use a more reasonable time.")
+            )
+    result["reason"] = duration.strip()
+    return result

forward/forward.py

@@ -2,12 +2,19 @@
 # from googletrans import Translator
 import discord
 import uuid
+from .converters import str_to_timedelta
+from datetime import datetime, timezone, timedelta
+import asyncio
+import logging
+
+from typing import Optional, Dict
 
 # trans = Translator()
 
+log = logging.getLogger("red.cogs.forward")
 
 class Forward(commands.Cog):
-    """Forward messages sent to the bot to the bot owner or in a specified channel."""
+    """Forward messages sent to the bot to the bot owner or a specified channel."""
 
     __version__ = "1.2.5"
 
@@ -27,6 +34,30 @@ def __init__(self, bot):
 
         self.welcome = {}
 
+        self._unblock_task = asyncio.create_task(self._handle_automatic_unblock())
+
+    def cog_unload(self):
+        self._unblock_task.cancel()
+
+    async def _handle_automatic_unblock(self):
+        """This is the core task creator and loop
+        for automatic unblocks
+        """
+        await self.bot.wait_until_red_ready()
+        while True:
+            try:
+                await self._process_expired_blocks()
+            except Exception:
+                log.error("Error checking bot unblocks", exc_info=True)
+            await asyncio.sleep(60) # 60 seconds
+
+    async def _process_expired_blocks(self):
+        async with self.config.blocked() as blocked:
+            for userid in list(blocked.keys()):
+                until = blocked[userid]
+                if datetime.now(timezone.utc).timestamp() > until:
+                    del blocked[userid]
+
     async def _destination(self, msg: str = None, embed: discord.Embed = None):
         await self.bot.wait_until_ready()
         channel = await self.config.destination()
@@ -240,28 +271,43 @@ async def tpm(self, ctx, language, user: discord.Member, *, message: str):
     @commands.command()
     @commands.guild_only()
     @checks.guildowner()
-    async def tblock(self, ctx, user: discord.Member):
-        """Blocks a member from sending dm
+    async def tblock(self, ctx, user: discord.Member, time: Optional[str] = "5d"):
+        """Blocks a member from sending DMs to the bot
         """
         async with self.config.blocked() as blocked:
             userid = str(user.id)
             if userid not in blocked:
-                blocked[userid] = True
-                await ctx.maybe_send_embed("Blocked <@{}> from send messages to the bot.".format(userid))
+                duration = str_to_timedelta(duration=time).get("duration")
+                duration = duration if type(duration) is timedelta else str_to_timedelta(duration="5d").get("duration")
+                blocked[userid] = (datetime.now(timezone.utc) + duration).timestamp() # until
+                await ctx.maybe_send_embed("Blocked <@{id}> from sending messages to the bot until <t:{until}>.".format(id=userid, until=int(blocked[userid])))
             else:
-                await ctx.maybe_send_embed("This user is already blocked.")
+                await ctx.maybe_send_embed("This user is already blocked until <t:{}>.".format(int(blocked[userid])))
 
     @commands.command()
     @commands.guild_only()
     @checks.guildowner()
     async def tunblock(self, ctx, user: discord.Member):
-        """Blocks a member from sending dm
+        """Unblocks a member from sending dms
         """
         async with self.config.blocked() as blocked:
             userid = str(user.id)
             if userid in blocked:
                 del blocked[userid]
-                await ctx.maybe_send_embed("Unblocked <@{}> from sending messages to the bot.".format(userid))
+                await ctx.maybe_send_embed("Unblocked <@{}> sending messages to the bot.".format(userid))
+            else:
+                await ctx.maybe_send_embed("This user is not blocked.")
+
+    @commands.command()
+    @commands.guild_only()
+    @checks.guildowner()
+    async def tcheckblock(self, ctx, user: discord.Member):
+        """Checks if a member is blocked from sending dms to the bot
+        """
+        async with self.config.blocked() as blocked:
+            userid = str(user.id)
+            if userid in blocked:
+                await ctx.maybe_send_embed("User <@{id}> is blocked until <t:{until}>.".format(id=userid, until=int(blocked[userid])))
             else:
                 await ctx.maybe_send_embed("This user is not blocked.")