diff --git a/bot/extensions/__init__.py b/bot/extensions/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bot/extensions/levelling/__init__.py b/bot/extensions/levelling/__init__.py new file mode 100644 index 00000000..0e10853a --- /dev/null +++ b/bot/extensions/levelling/__init__.py @@ -0,0 +1,9 @@ +from bot.core import DiscordBot + +from .commands import Levelling +from .events import LevelEvents + + +async def setup(bot: DiscordBot) -> None: + await bot.add_cog(Levelling(bot=bot)) + await bot.add_cog(LevelEvents(bot=bot)) diff --git a/bot/extensions/levelling/commands.py b/bot/extensions/levelling/commands.py new file mode 100644 index 00000000..ceafbe15 --- /dev/null +++ b/bot/extensions/levelling/commands.py @@ -0,0 +1,156 @@ +import datetime +import random + +import asyncpg.exceptions +import discord +from discord import app_commands +from discord.ext import commands + +from bot import core +from bot.models import IgnoredChannel, LevellingRole, Levels +from bot.models.custom_roles import CustomRoles + + +class Levelling(commands.Cog): + def __init__(self, bot): + self.bot = bot + self.ignored_channel = {} + self.required_xp = [0] + self.xp_boost = 1 + + async def cog_load(self): + for guild in self.bot.guilds: + data = await IgnoredChannel.list_by_guild(guild_id=guild.id) + for i in data: + if guild.id not in self.ignored_channel: + self.ignored_channel[guild.id] = [i.channel_id] + else: + self.ignored_channel[guild.id].append(i.channel_id) + + # Calculating required_XP for next level and storing in a list, list-index corresponds to the level + for lvl in range(101): + xp = 5 * (lvl**2) + (50 * lvl) + 100 + self.required_xp.append(xp + self.required_xp[-1]) + + @commands.Cog.listener() + async def on_message(self, message): + # Return if message was sent by Bot or sent in DMs + if message.author.bot or message.guild is None: + return + + # Check if message is sent in ignored channel + try: + if message.channel.id in self.ignored_channel[message.guild.id]: + return + except KeyError: + pass + + # Generate random XP to be added + xp = random.randint(5, 25) * self.xp_boost + # Add the XP and update the DB + data = await Levels.insert_by_guild(guild_id=message.guild.id, user_id=message.author.id, total_xp=xp) + self.bot.dispatch("xp_updated", data=data, member=message.author, required_xp=self.required_xp) + + @app_commands.command() + async def rank(self, interaction: core.InteractionType, member: discord.Member = None): + """Check the rank of another member or yourself""" + if member is None: + member = interaction.user + query = """WITH ordered_users AS ( + SELECT id,user_id,guild_id,total_xp, + ROW_NUMBER() OVER (ORDER BY levelling_users.total_xp DESC) AS rank + FROM levelling_users + WHERE guild_id = $2) + SELECT id,rank,total_xp,user_id,guild_id + FROM ordered_users WHERE ordered_users.user_id = $1;""" + data = await Levels.fetchrow(query, member.id, member.guild.id) + if data is None: + return await interaction.response.send_message("User Not ranked yet!", ephemeral=True) + for level, j in enumerate(self.required_xp): + if data.total_xp <= j: + embed = discord.Embed( + title=f"Rank: {data.rank}\nLevel: {level - 1}\nTotal XP:{data.total_xp}", + timestamp=datetime.datetime.utcnow(), + colour=discord.Colour.blurple(), + ) + embed.set_thumbnail(url=member.avatar) + return await interaction.response.send_message(embed=embed) + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def ignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel): + """Add the channel to the ignored channel list to not gain XP""" + await IgnoredChannel.insert_by_guild(channel.guild.id, channel.id) + await interaction.response.send_message(f"{channel} has been ignored from gaining XP.") + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def unignore_channel(self, interaction: core.InteractionType, channel: discord.TextChannel): + """Remove channel from ignored channel list""" + await IgnoredChannel.delete_by_guild(channel.guild.id, channel.id) + await interaction.response.send_message(f"{channel} has been removed from ignored channel list") + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def give_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member): + """Give XP to specific user""" + if xp <= 0: + return await interaction.response.send_message("XP can not be less than 0") + try: + data = await Levels.give_xp(member.guild.id, member.id, xp) + self.bot.dispatch("xp_updated", data=data, member=member, required_xp=self.required_xp) + await interaction.response.send_message(f"{xp} XP has been added to user {member}") + except asyncpg.exceptions.DataError: + return await interaction.response.send_message("Invalid XP provided") + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def remove_xp(self, interaction: core.InteractionType, xp: int, member: discord.Member): + """Remove XP from user""" + if xp <= 0: + return await interaction.response.send_message("XP can not be less than 0") + try: + data = await Levels.remove_xp(member.guild.id, member.id, xp) + self.bot.dispatch("xp_updated", data=data, member=member, required_xp=self.required_xp) + await interaction.response.send_message(f"{xp} XP has been removed from user {member}") + except asyncpg.exceptions.DataError: + return await interaction.response.send_message("Invalid XP provided") + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def levelling_rewards_add(self, interaction: core.InteractionType, role: discord.Role, level: int): + await CustomRoles.insert_by_guild(role.id, role.guild.id, role.name, str(role.colour)) + await LevellingRole.insert_by_guild(role.guild.id, role.id, level) + return await interaction.response.send_message("Levelling reward role added") + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def levelling_rewards_remove(self, interaction: core.InteractionType, role: discord.Role): + await LevellingRole.delete_by_guild(role.guild.id, role.id) + return await interaction.response.send_message("Levelling reward role removed") + + @app_commands.command() + async def levelling_rewards_list(self, interaction: core.InteractionType): + data = await LevellingRole.list_by_guild(interaction.guild.id) + res = "| {:<10} | {:<10} | {:<5} |".format("Guild ID", "Role ID", "Level") + res += "\n|" + "-" * 12 + "|" + "-" * 12 + "|" + "-" * 7 + "|" + + # Print data + for record in data: + res += "\n| {:<10} | {:<10} | {:<5} |".format(record["guild_id"], record["role_id"], record["level"]) + + await interaction.response.send_message(f"{res}") + + @app_commands.command() + @app_commands.checks.has_permissions(administrator=True) + async def xp_multiplier(self, interaction: core.InteractionType, multiplier: int): + """Increase XP gain per message""" + if multiplier <= 0 or multiplier > 5: + return await interaction.response.send_message("Invalid multiplier value.(Max. 5)") + + self.xp_boost = multiplier + return await interaction.response.send_message(f"XP multiplied by {multiplier}x.") + + +async def setup(bot: commands.Bot): + await bot.add_cog(Levelling(bot=bot)) diff --git a/bot/extensions/levelling/events.py b/bot/extensions/levelling/events.py new file mode 100644 index 00000000..80ab43bf --- /dev/null +++ b/bot/extensions/levelling/events.py @@ -0,0 +1,43 @@ +from bisect import bisect + +import discord +from discord.ext import commands + +from bot import core +from bot.models import LevellingRole, PersistentRole + + +class LevelEvents(commands.Cog): + """Events for Levelling in discord.""" + + def __init__(self, bot: core.DiscordBot): + self.bot = bot + + @commands.Cog.listener() + async def on_level_up(self, new_level, member: discord.Member): + """Add roles when user levels up""" + data = await LevellingRole.list_by_guild(guild_id=member.guild.id) + for i in range(len(data)): + # Adding roles when user levels up + if new_level >= data[i].level and member.guild.get_role(data[i].role_id) not in member.roles: + await member.add_roles(member.guild.get_role(data[i].role_id)) + await PersistentRole.insert_by_guild( + guild_id=member.guild.id, user_id=member.id, role_id=data[i].role_id + ) + # Removing roles when users level down using remove_xp command + elif new_level <= data[i].level and member.guild.get_role(data[i].role_id) in member.roles: + await member.remove_roles(member.guild.get_role(data[i].role_id)) + await PersistentRole.delete_by_guild(guild_id=member.guild.id, user_id=member.id) + + @commands.Cog.listener() + async def on_xp_updated(self, data, member: discord.Member, required_xp): + """Function to check if user's level has changed and trigger the event to assign the roles""" + # Calculating old and new level + try: + old_level = bisect(required_xp, data.old_total_xp) - 1 + new_level = bisect(required_xp, data.total_xp) - 1 + + if old_level != new_level: + self.bot.dispatch("level_up", new_level=new_level, member=member) + except AttributeError: + pass diff --git a/bot/extensions/persistent_roles/__init__.py b/bot/extensions/persistent_roles/__init__.py new file mode 100644 index 00000000..570498bc --- /dev/null +++ b/bot/extensions/persistent_roles/__init__.py @@ -0,0 +1,9 @@ +from bot.core import DiscordBot + +from .commands import PersistentRoles +from .events import PersistentEvents + + +async def setup(bot: DiscordBot) -> None: + await bot.add_cog(PersistentRoles(bot=bot)) + await bot.add_cog(PersistentEvents(bot=bot)) diff --git a/bot/extensions/persistent_roles/commands.py b/bot/extensions/persistent_roles/commands.py new file mode 100644 index 00000000..72769c53 --- /dev/null +++ b/bot/extensions/persistent_roles/commands.py @@ -0,0 +1,10 @@ +from discord.ext import commands + + +class PersistentRoles(commands.Cog): + def __init__(self, bot): + self.bot = bot + + +async def setup(bot: commands.Bot): + await bot.add_cog(PersistentRoles(bot=bot)) diff --git a/bot/extensions/persistent_roles/events.py b/bot/extensions/persistent_roles/events.py new file mode 100644 index 00000000..3e3bf07c --- /dev/null +++ b/bot/extensions/persistent_roles/events.py @@ -0,0 +1,24 @@ +from discord.ext import commands + +from bot import core +from bot.models import PersistentRole + + +class PersistentEvents(commands.Cog): + """Events for Persisted roles.""" + + def __init__(self, bot: core.DiscordBot): + self.bot = bot + + @commands.Cog.listener() + async def on_member_join(self, member): + """Add the persisted role to users if any on member join""" + # Get the data + data = await PersistentRole.list_by_guild(member.guild.id, member.id) + # Return if data for specified user and guild does not exist + if data is None: + return + # Add the roles to the user if data exists + else: + for i in range(len(data)): + await member.add_roles(member.guild.get_role(data[i].role_id)) diff --git a/bot/models/__init__.py b/bot/models/__init__.py index ab0e68f6..c80a4bfa 100644 --- a/bot/models/__init__.py +++ b/bot/models/__init__.py @@ -1,15 +1,23 @@ from .gconfig import FilterConfig +from .levelling_ignored_channels import IgnoredChannel +from .levelling_roles import LevellingRole +from .levelling_users import Levels from .message import Message from .model import Model +from .persisted_role import PersistentRole from .rep import Rep from .tag import Tag from .user import User -__all__ = ( # Fixes F401 +__all__ = ( Model, FilterConfig, Message, Rep, Tag, User, -) + Levels, + PersistentRole, + IgnoredChannel, + LevellingRole, +) # Fixes F401 diff --git a/bot/models/custom_roles.py b/bot/models/custom_roles.py new file mode 100644 index 00000000..a764ea66 --- /dev/null +++ b/bot/models/custom_roles.py @@ -0,0 +1,20 @@ +from .model import Model + + +class CustomRoles(Model): + id: int + role_id: int + guild_id: int + name: str + color: str + + @classmethod + async def insert_by_guild(cls, role_id: int, guild_id: int, name: str, color: str): + query = """INSERT INTO custom_roles (role_id, guild_id, name, color) + VALUES($1, $2, $3, $4)""" + await cls.execute(query, role_id, guild_id, name, color) + + @classmethod + async def delete_by_guild(cls, guild_id: int, role_id: int): + query = """DELETE FROM custom_roles WHERE guild_id = $1 and role_id = $2""" + await cls.execute(query, guild_id, role_id) diff --git a/bot/models/levelling_ignored_channels.py b/bot/models/levelling_ignored_channels.py new file mode 100644 index 00000000..931ac7e4 --- /dev/null +++ b/bot/models/levelling_ignored_channels.py @@ -0,0 +1,24 @@ +from .model import Model + + +class IgnoredChannel(Model): + id: int + guild_id: int + channel_id: int + + @classmethod + async def list_by_guild(cls, guild_id: int): + query = """SELECT * FROM levelling_ignored_channels WHERE guild_id = $1""" + return await cls.fetch(query, guild_id) + + @classmethod + async def insert_by_guild(cls, guild_id: int, channel_id: int): + query = """INSERT INTO levelling_ignored_channels (guild_id,channel_id) VALUES($1, $2) + ON CONFLICT (guild_id,channel_id) DO NOTHING + RETURNING id, guild_id, channel_id""" + return await cls.fetch(query, guild_id, channel_id) + + @classmethod + async def delete_by_guild(cls, guild_id: int, channel_id: int): + query = """DELETE FROM levelling_ignored_channels WHERE guild_id= $1 and channel_id = $2""" + return await cls.fetch(query, guild_id, channel_id) diff --git a/bot/models/levelling_roles.py b/bot/models/levelling_roles.py new file mode 100644 index 00000000..83340353 --- /dev/null +++ b/bot/models/levelling_roles.py @@ -0,0 +1,26 @@ +from .model import Model + + +class LevellingRole(Model): + id: int + guild_id: int + role_id: int + level: int + + @classmethod + async def list_by_guild(cls, guild_id: int): + query = """SELECT * FROM levelling_roles WHERE guild_id=$1""" + return await cls.fetch(query, guild_id) + + @classmethod + async def insert_by_guild(cls, guild_id: int, role_id: int, level: int): + query = """INSERT INTO levelling_roles (guild_id, role_id, level) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, role_id) + DO NOTHING""" + await cls.execute(query, guild_id, role_id, level) + + @classmethod + async def delete_by_guild(cls, guild_id: int, role_id: int): + query = """DELETE FROM levelling_roles WHERE guild_id = $1 and role_id = $2""" + await cls.execute(query, guild_id, role_id) diff --git a/bot/models/levelling_users.py b/bot/models/levelling_users.py new file mode 100644 index 00000000..fa66f1a6 --- /dev/null +++ b/bot/models/levelling_users.py @@ -0,0 +1,48 @@ +from typing import Optional + +from .model import Model + + +class Levels(Model): + id: int + rank: Optional[int] + guild_id: int + user_id: int + old_total_xp: Optional[int] + total_xp: int + + @classmethod + async def insert_by_guild(cls, guild_id: int, user_id: int, total_xp: int): + query = """INSERT INTO levelling_users (guild_id, user_id, total_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id) + DO UPDATE + SET total_xp = levelling_users.total_xp + $3, last_msg = CURRENT_TIMESTAMP + WHERE levelling_users.guild_id = $1 + AND levelling_users.user_id = $2 + AND EXTRACT(EPOCH FROM (NOW() - levelling_users.last_msg)) > 60 + RETURNING total_xp - $3 AS old_total_xp, total_xp, guild_id, user_id, id; +""" + return await cls.fetchrow(query, guild_id, user_id, total_xp) + + @classmethod + async def give_xp(cls, guild_id: int, user_id: int, total_xp: int): + query = """INSERT INTO levelling_users (guild_id, user_id, total_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id) + DO UPDATE SET total_xp = levelling_users.total_xp + $3 + RETURNING total_xp - $3 AS old_total_xp, total_xp, guild_id, user_id, id; + """ + + return await cls.fetchrow(query, guild_id, user_id, total_xp) + + @classmethod + async def remove_xp(cls, guild_id: int, user_id: int, total_xp: int): + query = """INSERT INTO levelling_users (guild_id, user_id, total_xp) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id) + DO UPDATE SET total_xp = levelling_users.total_xp - $3 + RETURNING total_xp + $3 AS old_total_xp, total_xp, guild_id, user_id, id; + """ + + return await cls.fetchrow(query, guild_id, user_id, total_xp) diff --git a/bot/models/migrations/002_down__levelling_db.sql b/bot/models/migrations/002_down__levelling_db.sql new file mode 100644 index 00000000..5b043205 --- /dev/null +++ b/bot/models/migrations/002_down__levelling_db.sql @@ -0,0 +1 @@ +DROP TABLE levels, roles, persistent_roles, ignored_channels, levelling_roles; diff --git a/bot/models/migrations/002_up__levelling_db.sql b/bot/models/migrations/002_up__levelling_db.sql new file mode 100644 index 00000000..5a59b781 --- /dev/null +++ b/bot/models/migrations/002_up__levelling_db.sql @@ -0,0 +1,48 @@ +CREATE TABLE IF NOT EXISTS levelling_users +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + user_id BIGINT, + total_xp INT DEFAULT 0, + last_msg TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + CONSTRAINT levelling_roles_guild_id_and_user_id_key UNIQUE (guild_id, user_id) +); + +CREATE TABLE IF NOT EXISTS custom_roles +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + role_id BIGINT UNIQUE, + name VARCHAR, + color VARCHAR +); + +CREATE TABLE IF NOT EXISTS persisted_roles +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + user_id BIGINT, + role_id BIGINT, + CONSTRAINT persisted_roles_role_id_fkey FOREIGN KEY(role_id) REFERENCES custom_roles(role_id), + CONSTRAINT persisted_roles_guild_id_and_user_id_and_role_id_key UNIQUE (guild_id, user_id, role_id) +); + + +CREATE TABLE IF NOT EXISTS levelling_ignored_channels +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + channel_id BIGINT, + CONSTRAINT levelling_ignored_channels_guild_id_and_channel_id_key UNIQUE (guild_id, channel_id) + +); + +CREATE TABLE IF NOT EXISTS levelling_roles +( + id SERIAL PRIMARY KEY, + guild_id BIGINT, + role_id BIGINT, + level INT, + CONSTRAINT levelling_roles_role_id_fkey FOREIGN KEY(role_id) REFERENCES custom_roles(role_id), + CONSTRAINT levelling_roles_guild_id_and_role_id_key UNIQUE (guild_id, role_id) +); diff --git a/bot/models/persisted_role.py b/bot/models/persisted_role.py new file mode 100644 index 00000000..3c28b0d6 --- /dev/null +++ b/bot/models/persisted_role.py @@ -0,0 +1,29 @@ +from typing import Optional + +from .model import Model + + +class PersistentRole(Model): + id: int + guild_id: int + user_id: int + role_id: Optional[int] + + @classmethod + async def list_by_guild(cls, guild_id: int, user_id: int): + query = """SELECT * FROM persisted_roles WHERE guild_id = $1 and user_id = $2""" + return await cls.fetch(query, guild_id, user_id) + + @classmethod + async def insert_by_guild(cls, guild_id: int, user_id: int, role_id: int): + query = """INSERT INTO persisted_roles (guild_id, user_id, role_id) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, user_id, role_id) + DO NOTHING""" + + await cls.execute(query, guild_id, user_id, role_id) + + @classmethod + async def delete_by_guild(cls, guild_id: int, user_id: int): + query = """DELETE FROM persisted_roles WHERE guild_id = $1 and user_id = $2""" + return await cls.execute(query, guild_id, user_id) diff --git a/cli.py b/cli.py index 5e7c8948..2f7a70e3 100644 --- a/cli.py +++ b/cli.py @@ -139,6 +139,8 @@ async def main(ctx): "bot.cogs.clashofcode", "bot.cogs.roles", "bot.cogs.poll", + "bot.extensions.levelling", + "bot.extensions.persistent_roles", ) intents = discord.Intents.all()