Work on rob

database/crud/currency.py

@@ -47,6 +47,9 @@ async def invest(session: AsyncSession, user_id: int, amount: Union[str, int]) -
     if amount == "all":
         amount = bank.dinks
 
+    if bank.dinks <= 0:
+        return 0
+
     # Don't allow investing more dinks than you own
     amount = min(bank.dinks, int(amount))
 
@@ -158,11 +161,14 @@ async def upgrade_rob(session: AsyncSession, user_id: int) -> int:
 async def gamble_dinks(
     session: AsyncSession, user_id: int, amount: Union[str, int], payout_factor: int, won: bool
 ) -> int:
-    """Gamble some of your Dinks"""
+    """Gamble some of your Didier Dinks"""
     bank = await get_bank(session, user_id)
     if amount == "all":
         amount = bank.dinks
 
+    if bank.dinks <= 0:
+        return 0
+
     amount = min(int(amount), bank.dinks)
 
     sign = 1 if won else -1
@@ -170,3 +176,16 @@ async def gamble_dinks(
     await add_dinks(session, user_id, sign * amount * factor)
 
     return amount * factor
+
+
+async def rob(session: AsyncSession, amount: int, robber_id: int, robbed_id: int):
+    """Rob another user's Didier Dinks"""
+    robber = await get_bank(session, robber_id)
+    robbed = await get_bank(session, robbed_id)
+
+    robber.dinks += amount
+    robbed.dinks -= amount
+
+    session.add(robber)
+    session.add(robbed)
+    await session.commit()

database/crud/jail.py

@@ -0,0 +1,28 @@
+from datetime import datetime
+from typing import Optional
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from database.schemas import Jail
+
+__all__ = ["get_jail", "get_user_jail", "imprison"]
+
+
+async def get_jail(session: AsyncSession) -> list[Jail]:
+    """Get the entire Didier Jail"""
+    statement = select(Jail)
+    return list((await session.execute(statement)).scalars().all())
+
+
+async def get_user_jail(session: AsyncSession, user_id: int) -> Optional[Jail]:
+    """Check how long a given user is still in jail for"""
+    statement = select(Jail).where(Jail.user_id == user_id)
+    return (await session.execute(statement)).scalar_one_or_none()
+
+
+async def imprison(session: AsyncSession, user_id: int, until: datetime):
+    """Put a user in Didier Jail"""
+    jail = Jail(user_id=user_id, until=until)
+    session.add(jail)
+    await session.commit()

database/schemas.py

@@ -23,6 +23,7 @@ __all__ = [
     "Event",
     "FreeGame",
     "GitHubLink",
+    "Jail",
     "Link",
     "MemeTemplate",
     "NightlyData",
@@ -199,6 +200,18 @@ class GitHubLink(Base):
     user: Mapped[User] = relationship(back_populates="github_links", uselist=False, lazy="selectin")
 
 
+class Jail(Base):
+    """A user sitting in Didier Jail"""
+
+    __tablename__ = "jail"
+
+    jail_entry_i: Mapped[int] = mapped_column(primary_key=True)
+    user_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("users.user_id"))
+    until: Mapped[datetime] = mapped_column(nullable=False)
+
+    user: Mapped[User] = relationship(back_populates="jail", uselist=False, lazy="selectin")
+
+
 class Link(Base):
     """Useful links that go useful places"""
 
@@ -328,6 +341,9 @@ class User(Base):
     github_links: Mapped[List[GitHubLink]] = relationship(
         back_populates="user", uselist=True, lazy="selectin", cascade="all, delete-orphan"
     )
+    jail: Mapped[Optional[Jail]] = relationship(
+        back_populates="user", uselist=False, lazy="selectin", cascade="all, delete-orphan"
+    )
     nightly_data: Mapped[NightlyData] = relationship(
         back_populates="user", uselist=False, lazy="selectin", cascade="all, delete-orphan"
     )

database/utils/math/currency.py

@@ -1,6 +1,14 @@
 import math
 
-__all__ = ["capacity_upgrade_price", "interest_upgrade_price", "rob_upgrade_price"]
+__all__ = [
+    "capacity_upgrade_price",
+    "interest_upgrade_price",
+    "rob_upgrade_price",
+    "jail_chance",
+    "jail_time",
+    "rob_amount",
+    "rob_chance",
+]
 
 
 def interest_upgrade_price(level: int) -> int:
@@ -25,3 +33,34 @@ def rob_upgrade_price(level: int) -> int:
     growth_rate = 1.9
 
     return math.floor(base_cost * (growth_rate**level))
+
+
+def jail_chance(level: int) -> float:
+    """Calculate the chance that you'll end up in jail"""
+    base_chance = 0.35
+    growth_rate = 1.15
+
+    return max(0.0, base_chance - (growth_rate**level))
+
+
+def jail_time(level: int) -> int:
+    """Calculate the time in hours you'll have to spend in jail"""
+    base_hours = 2
+    growth_rate = 1.27
+
+    return math.floor(base_hours + growth_rate**level)
+
+
+def rob_amount(level: int) -> int:
+    """Calculate the maximum amount of Didier Dinks that you can rob"""
+    base_amount = 250
+    growth_rate = 1.4
+
+    return math.floor(base_amount * (growth_rate**level))
+
+
+def rob_chance(level: int) -> float:
+    """Calculate the chances of a robbing attempt being successful"""
+    base_chance = 0.25
+
+    return base_chance + 2.1 * level

didier/cogs/currency.py

@@ -1,17 +1,25 @@
 # flake8: noqa: E800
+import asyncio
+import math
+import random
 import typing
 
 import discord
 from discord.ext import commands
 
+import settings
 from database.crud import currency as crud
+from database.crud.jail import get_user_jail
 from database.exceptions.currency import DoubleNightly, NotEnoughDinks
 from database.utils.math.currency import (
     capacity_upgrade_price,
     interest_upgrade_price,
+    rob_amount,
+    rob_chance,
     rob_upgrade_price,
 )
 from didier import Didier
+from didier.utils.discord import colours
 from didier.utils.discord.checks import is_owner
 from didier.utils.discord.converters import abbreviated_number
 from didier.utils.types.string import pluralize
@@ -21,10 +29,12 @@ class Currency(commands.Cog):
     """Everything Dinks-related."""
 
     client: Didier
+    _rob_lock: asyncio.Lock
 
     def __init__(self, client: Didier):
         super().__init__()
         self.client = client
+        self._rob_lock = asyncio.Lock()
 
     @commands.command(name="award")  # type: ignore[arg-type]
     @commands.check(is_owner)
@@ -61,9 +71,9 @@ class Currency(commands.Cog):
 
         await ctx.reply(embed=embed, mention_author=False)
 
-    # @bank.group(  # type: ignore[arg-type]
+    @bank.group(  # type: ignore[arg-type]
-    #     name="upgrade", aliases=["u", "upgrades"], case_insensitive=True, invoke_without_command=True
+        name="upgrade", aliases=["u", "upgrades"], case_insensitive=True, invoke_without_command=True
-    # )
+    )
     async def bank_upgrades(self, ctx: commands.Context):
         """List the upgrades you can buy & their prices."""
         async with self.client.postgres_session as session:
@@ -83,7 +93,7 @@ class Currency(commands.Cog):
 
         await ctx.reply(embed=embed, mention_author=False)
 
-    # @bank_upgrades.command(name="capacity", aliases=["c"])  # type: ignore[arg-type]
+    @bank_upgrades.command(name="capacity", aliases=["c"])  # type: ignore[arg-type]
     async def bank_upgrade_capacity(self, ctx: commands.Context):
         """Upgrade the capacity level of your bank."""
         async with self.client.postgres_session as session:
@@ -94,7 +104,7 @@ class Currency(commands.Cog):
                 await ctx.reply("You don't have enough Didier Dinks to do this.", mention_author=False)
                 await self.client.reject_message(ctx.message)
 
-    # @bank_upgrades.command(name="interest", aliases=["i"])  # type: ignore[arg-type]
+    @bank_upgrades.command(name="interest", aliases=["i"])  # type: ignore[arg-type]
     async def bank_upgrade_interest(self, ctx: commands.Context):
         """Upgrade the interest level of your bank."""
         async with self.client.postgres_session as session:
@@ -105,7 +115,7 @@ class Currency(commands.Cog):
                 await ctx.reply("You don't have enough Didier Dinks to do this.", mention_author=False)
                 await self.client.reject_message(ctx.message)
 
-    # @bank_upgrades.command(name="rob", aliases=["r"])  # type: ignore[arg-type]
+    @bank_upgrades.command(name="rob", aliases=["r"])  # type: ignore[arg-type]
     async def bank_upgrade_rob(self, ctx: commands.Context):
         """Upgrade the rob level of your bank."""
         async with self.client.postgres_session as session:
@@ -182,6 +192,70 @@ class Currency(commands.Cog):
                     "You've already claimed your Didier Nightly today.", mention_author=False, ephemeral=True
                 )
 
+    @commands.hybrid_command(name="rob")  # type: ignore[arg-type]
+    @commands.cooldown(rate=1, per=10 * 60.0, type=commands.BucketType.user)
+    @commands.guild_only()
+    async def rob(self, ctx: commands.Context, member: discord.Member):
+        """Attempt to rob another user of their Dinks"""
+        if member == ctx.author:
+            return await ctx.reply("You can't rob yourself.", mention_author=False, ephemeral=True)
+
+        if member == self.client.user:
+            return await ctx.reply(settings.DISCORD_BOOS_REACT, mention_author=False)
+
+        if member.bot:
+            return await ctx.reply("You can't rob bots.", mention_author=False, ephemeral=True)
+
+        # Use a Lock for robbing to avoid race conditions when robbing the same person twice
+        # This would cause undefined behaviour
+        async with ctx.typing(), self._rob_lock, self.client.postgres_session as session:
+            robber = await crud.get_bank(session, ctx.author.id)
+            robbed = await crud.get_bank(session, member.id)
+
+            if robber.dinks <= 0:
+                return await ctx.reply(
+                    "You can't rob without dinks. Just stop being poor lol", mention_author=False, ephemeral=True
+                )
+
+            if robbed.dinks <= 0:
+                return await ctx.reply(
+                    f"{member.display_name} doesn't have any dinks to rob.", mention_author=False, ephemeral=True
+                )
+
+            rob_roll = random.random()
+            success_chance = rob_chance(robber.rob_level)
+            success = rob_roll <= success_chance
+
+            if success:
+                max_rob_amount = random.uniform(0.15, 1.0) * rob_amount(robber.rob_level)
+                robbed_amount = min(robbed.dinks, math.floor(max_rob_amount))
+                await crud.rob(session, robbed_amount, ctx.author.id, member.id)
+                return await ctx.reply(
+                    f"{ctx.author.display_name} has robbed **{robbed_amount}** Didier Dinks from {member.display_name}!",
+                    mention_author=False,
+                )
+
+    @commands.hybrid_command(name="jail")
+    async def jail(self, ctx: commands.Context):
+        """Check how long you're still in jail for"""
+        async with self.client.postgres_session as session:
+            entry = await get_user_jail(session, ctx.author.id)
+
+        if entry is None:
+            embed = discord.Embed(
+                title="Didier Jail", colour=colours.error_red(), description="You're not currently in jail."
+            )
+
+            return await ctx.reply(embed=embed, mention_author=False, ephemeral=True)
+
+        embed = discord.Embed(
+            title="Didier Jail",
+            colour=colours.jail_gray(),
+            description=f"You will be released <t:{entry.until.timestamp()}:R>.",
+        )
+
+        return await ctx.reply(embed=embed, mention_author=False)
+
 
 async def setup(client: Didier):
     """Load the cog"""

didier/didier.py

@@ -310,6 +310,14 @@ class Didier(commands.Bot):
         ):
             return await ctx.reply(str(exception.original), mention_author=False)
 
+        if isinstance(exception, (discord.app_commands.CommandOnCooldown, commands.CommandOnCooldown)):
+            return await ctx.reply(
+                f"You have to wait another {exception.retry_after} seconds before "
+                f"attempting to use this command again.",
+                ephemeral=True,
+                delete_after=10,
+            )
+
         if isinstance(exception, commands.MessageNotFound):
             return await ctx.reply("This message could not be found.", ephemeral=True, delete_after=10)
 
@@ -332,6 +340,12 @@ class Didier(commands.Bot):
         ):
             return await ctx.reply("Invalid arguments.", ephemeral=True, delete_after=10)
 
+        if isinstance(exception, commands.NoPrivateMessage):
+            if ctx.command.name == "rob":
+                return await ctx.reply("Robbing in DM? That's low, even for you.", ephemeral=True)
+
+            return await ctx.reply("This command cannot be used in DMs.", ephemeral=True)
+
         # Print everything that we care about to the logs/stderr
         await super().on_command_error(ctx, exception)
 

didier/utils/discord/colours.py

@@ -7,6 +7,7 @@ __all__ = [
     "ghent_university_blue",
     "ghent_university_yellow",
     "google_blue",
+    "jail_gray",
     "steam_blue",
     "urban_dictionary_green",
     "xkcd_blue",
@@ -41,6 +42,10 @@ def google_blue() -> discord.Colour:
     return discord.Colour.from_rgb(66, 133, 244)
 
 
+def jail_gray() -> discord.Colour:
+    return discord.Colour.from_rgb(85, 85, 85)
+
+
 def steam_blue() -> discord.Colour:
     return discord.Colour.from_rgb(102, 192, 244)