From fa07c6b2e9ab345f27748ed30fe7e65e17e8f5ae Mon Sep 17 00:00:00 2001
From: "Nyeki.py" <84998222+Nekidev@users.noreply.github.com>
Date: Fri, 18 Oct 2024 04:38:00 -0300
Subject: [PATCH] feat(bot): :sparkles: Make ban durations stack.

---
 payphone/bot.py                               | 23 ++++++++++++++++---
 payphone/cogs/mods/cog.py                     | 11 +++++----
 payphone/cogs/mods/modals.py                  |  7 ++++--
 .../models/7_20241018043647_update.py         | 11 +++++++++
 payphone/db/models.py                         |  2 +-
 5 files changed, 44 insertions(+), 10 deletions(-)
 create mode 100644 payphone/db/migrations/models/7_20241018043647_update.py

diff --git a/payphone/bot.py b/payphone/bot.py
index 68b0127..c355c11 100644
--- a/payphone/bot.py
+++ b/payphone/bot.py
@@ -1,6 +1,6 @@
 from typing import Literal
 
-from datetime import datetime, timezone
+from datetime import datetime, timezone, timedelta
 
 import discord
 from discord import abc
@@ -147,7 +147,7 @@ class Payphone(commands.AutoShardedBot):
     async def ban(
         self,
         user: discord.User | discord.Member | models.User | int,
-        until: datetime | None,
+        duration: timedelta | None,
         reason: str,
         moderator: discord.User | discord.Member | models.User | int | None = None,
     ) -> models.Ban:
@@ -155,7 +155,7 @@ class Payphone(commands.AutoShardedBot):
 
         Args:
             user (discord.User | discord.Member | models.User | int): The user to ban or their ID.
-            until (datetime | None): The expiry date, or None if it does not expire.
+            duration (timedelta | None): The expiry date, or None if it does not expire.
             reason (str): The reason of the ban.
             moderator (discord.User | discord.Member | models.User | int | None, optional): The
               moderator who performed the action.
@@ -173,6 +173,23 @@ class Payphone(commands.AutoShardedBot):
             else moderator
         )
 
+        until = datetime.now(timezone.utc) + duration if duration is not None else None
+
+        if duration is not None:
+            last_ban = (
+                await models.Ban.filter(user_id=user_id)
+                .order_by("-created_at")
+                .limit(1)
+                .first()
+            )
+
+            if last_ban is not None:
+                until = (
+                    duration + last_ban.expires_at
+                    if last_ban.expires_at is not None
+                    else None
+                )
+
         row = await models.Ban.create(
             user_id=user_id,
             reason=reason,
diff --git a/payphone/cogs/mods/cog.py b/payphone/cogs/mods/cog.py
index 4f849a9..4f7b703 100644
--- a/payphone/cogs/mods/cog.py
+++ b/payphone/cogs/mods/cog.py
@@ -31,7 +31,7 @@ class Mods(base.Cog):
         self,
         ctx: commands.Context,
         user: converters.DiscordUserID,
-        time: converters.Timedelta | None = None,
+        duration: converters.Timedelta | None = None,
         *,
         reason: str | None = "",
     ) -> None:
@@ -40,7 +40,7 @@ class Mods(base.Cog):
 
         Options:
         - `<user>`: The user to ban or their ID.
-        - `[<time>]`: The amount of time to ban the user for. For example, `1d 5h`.
+        - `[<duration>]`: The amount of time to ban the user for. For example, `1d 5h`.
         - `[<reason>]`: The reason of their ban.
         """
 
@@ -50,7 +50,7 @@ class Mods(base.Cog):
 
         row = await self.bot.ban(
             user_id,
-            (datetime.now(timezone.utc) + time) if time is not None else None,
+            duration,
             reason,
             ctx.author,
         )
@@ -58,7 +58,10 @@ class Mods(base.Cog):
         message = f"✅ {f"**{profile.name}**" if profile is not None else f"`ID: {user_id}`"} was banned"
 
         if row.expires_at:
-            message += f" until <t:{int(row.expires_at.timestamp())}:F> (<t:{int(row.expires_at.timestamp())}:R>)"
+            message += f"{f" for `{fmt.time_to_text(duration)}`" if duration is not None else ""} until <t:{int(row.expires_at.timestamp())}:F> (<t:{int(row.expires_at.timestamp())}:R>)"
+        
+        else:
+            message += " permanently"
 
         message += "."
 
diff --git a/payphone/cogs/mods/modals.py b/payphone/cogs/mods/modals.py
index 288db5a..056eb19 100644
--- a/payphone/cogs/mods/modals.py
+++ b/payphone/cogs/mods/modals.py
@@ -98,7 +98,7 @@ class BanUserModal(discord.ui.Modal):
 
         ban = await self.bot.ban(
             self.profile.user_id,
-            datetime.now(timezone.utc) + duration if duration is not None else None,
+            duration,
             self.children[0].value,
             interaction.user.id,
         )
@@ -120,7 +120,10 @@ class BanUserModal(discord.ui.Modal):
         message = f"✅ **{self.profile.name}** (User ID: {self.profile.user_id}) was banned by {interaction.user.mention}"
 
         if ban.expires_at is not None:
-            message += f" until <t:{int(ban.expires_at.timestamp())}:F> (<t:{int(ban.expires_at.timestamp())}:R>)"
+            message += f"{f" for `{formatting.time_to_text(duration)}`" if duration is not None else ""} until <t:{int(ban.expires_at.timestamp())}:F> (<t:{int(ban.expires_at.timestamp())}:R>)"
+        
+        else:
+            message += " permanently"
 
         message += "."
 
diff --git a/payphone/db/migrations/models/7_20241018043647_update.py b/payphone/db/migrations/models/7_20241018043647_update.py
new file mode 100644
index 0000000..b313887
--- /dev/null
+++ b/payphone/db/migrations/models/7_20241018043647_update.py
@@ -0,0 +1,11 @@
+from tortoise import BaseDBAsyncClient
+
+
+async def upgrade(db: BaseDBAsyncClient) -> str:
+    return """
+        CREATE INDEX "idx_ban_created_490013" ON "ban" ("created_at");"""
+
+
+async def downgrade(db: BaseDBAsyncClient) -> str:
+    return """
+        DROP INDEX "idx_ban_created_490013";"""
diff --git a/payphone/db/models.py b/payphone/db/models.py
index 1f3ad8f..917d935 100644
--- a/payphone/db/models.py
+++ b/payphone/db/models.py
@@ -346,7 +346,7 @@ class Ban(models.Model):
     banned_by: fields.ForeignKeyRelation["User"] = fields.ForeignKeyField(
         "models.User", related_name="bans_created", null=True
     )
-    created_at: datetime = fields.DatetimeField(auto_now_add=True)
+    created_at: datetime = fields.DatetimeField(auto_now_add=True, db_index=True)
     expires_at: datetime | None = fields.DatetimeField(null=True)
 
     class Meta: