diff --git a/.example.env b/.example.env index aa11f8b..46c41ae 100644 --- a/.example.env +++ b/.example.env @@ -7,3 +7,4 @@ GEMINI_API_KEY="GEMINI_API_KEY" REQUESTS_PER_MINUTE=3 LIMIT_WINDOW=60 LOG_CHANNEL_ID=LOG_CHANNEL_ID +ADMIN_USERS=User1,User2,User3 diff --git a/.github/workflows/production.yml b/.github/workflows/production.yml index 703fa5c..89667a2 100644 --- a/.github/workflows/production.yml +++ b/.github/workflows/production.yml @@ -85,6 +85,7 @@ jobs: REQUESTS_PER_MINUTE: ${{ secrets.REQUESTS_PER_MINUTE }} LIMIT_WINDOW: ${{ secrets.LIMIT_WINDOW }} LOG_CHANNEL_ID: ${{ secrets.LOG_CHANNEL_ID }} + ADMIN_USERS: ${{ secrets.ADMIN_USERS }} run: | echo "$KEY" > private_key && chmod 600 private_key ssh -v -o StrictHostKeyChecking=no -i private_key ${USER}@${HOSTNAME} ' @@ -100,6 +101,7 @@ jobs: echo REQUESTS_PER_MINUTE=${{ secrets.REQUESTS_PER_MINUTE }} >> .env echo LIMIT_WINDOW=${{ secrets.LIMIT_WINDOW }} >> .env echo LOG_CHANNEL_ID=${{ secrets.LOG_CHANNEL_ID }} >> .env + echo ADMIN_USERS=${{ secrets.ADMIN_USERS }} >> .env docker load -i duckbot.tar.gz docker compose up -d ' diff --git a/src/commands/admin_commands.py b/src/commands/admin_commands.py new file mode 100644 index 0000000..60b37f2 --- /dev/null +++ b/src/commands/admin_commands.py @@ -0,0 +1,128 @@ +import os +import logging +from discord import app_commands, Interaction, Embed +from models.databases.admin_settings_db import AdminSettingsDB + +# Retrieve the list of admin usernames from the .env file +ADMIN_USERS = os.getenv("ADMIN_USERS", "").split(",") + + +class AdminCommands(app_commands.Group): + def __init__(self, gemini_bot): + super().__init__(name="admin", description="Admin commands for DuckBot setup.") + + # Initialise the database + self.settings_db = AdminSettingsDB() + + # Add subgroups to the main admin group + self.set = SetSubGroup(self.check_admin, self.settings_db) + self.reset = ResetSubGroup(self.check_admin, gemini_bot) + + # Register subgroups + self.add_command(self.set) + self.add_command(self.reset) + + async def check_admin(self, interaction: Interaction) -> bool: + user_name = interaction.user.name + logging.info(f"Checking admin status for user: {user_name}") + + if user_name in ADMIN_USERS: + logging.info(f"User {user_name} is authorised.") + return True + else: + await interaction.response.send_message( + "You don't have permission to execute that command.", ephemeral=True + ) + logging.warning(f"User {user_name} is not authorised.") + return False + + @app_commands.command( + name="log-variables", description="Display all current environment variables." + ) + async def log_info(self, interaction: Interaction): + """Command to log and display all relevant environment variables.""" + if not await self.check_admin(interaction): + return + + # Get values from database instead of env + guild_id = self.settings_db.get_setting("GUILD_ID") or "Not Set" + skullboard_channel_id = ( + self.settings_db.get_setting("SKULLBOARD_CHANNEL_ID") or "Not Set" + ) + required_reactions = ( + self.settings_db.get_setting("REQUIRED_REACTIONS") or "Not Set" + ) + + embed = Embed(title="Current Settings", color=0x00FF00) + embed.add_field(name="Guild ID", value=f"`{guild_id}`", inline=False) + embed.add_field( + name="Skullboard Channel ID", + value=f"`{skullboard_channel_id}`", + inline=False, + ) + embed.add_field( + name="Required Reactions", value=f"`{required_reactions}`", inline=False + ) + + await interaction.response.send_message(embed=embed, ephemeral=True) + + +class SetSubGroup(app_commands.Group): + def __init__(self, check_admin, settings_db): + super().__init__( + name="set", description="Set configuration values for DuckBot." + ) + self.check_admin = check_admin + self.settings_db = settings_db + + @app_commands.command(name="guild-id", description="Set the guild ID for DuckBot.") + async def set_guild_id(self, interaction: Interaction, guild_id: str): + if not await self.check_admin(interaction): + return + self.settings_db.set_setting("GUILD_ID", guild_id) + await interaction.response.send_message( + f"Guild ID set to {guild_id}.", ephemeral=True + ) + + @app_commands.command( + name="skullboard-channel-id", description="Set the Skullboard channel ID." + ) + async def set_skullboard_channel_id( + self, interaction: Interaction, channel_id: str + ): + if not await self.check_admin(interaction): + return + self.settings_db.set_setting("SKULLBOARD_CHANNEL_ID", channel_id) + await interaction.response.send_message( + f"Skullboard channel ID set to {channel_id}.", ephemeral=True + ) + + @app_commands.command( + name="required-reactions", description="Set required reactions for Skullboard." + ) + async def set_required_reactions(self, interaction: Interaction, reactions: int): + if not await self.check_admin(interaction): + return + self.settings_db.set_setting("REQUIRED_REACTIONS", str(reactions)) + await interaction.response.send_message( + f"Required reactions set to {reactions}.", ephemeral=True + ) + + +class ResetSubGroup(app_commands.Group): + def __init__(self, check_admin, gemini_bot): + super().__init__(name="reset", description="Reset specific DuckBot settings.") + self.check_admin = check_admin + self.gemini_bot = gemini_bot + + @app_commands.command(name="chat-history", description="Reset Gemini chat history.") + async def reset_chat_history(self, interaction: Interaction): + if not await self.check_admin(interaction): + return + + # Call the method to reset Gemini chat history + self.gemini_bot.clear_chat_history() + + await interaction.response.send_message( + "Gemini chat history has been reset.", ephemeral=True + ) diff --git a/src/main.py b/src/main.py index 3e4de0a..7f3826c 100644 --- a/src/main.py +++ b/src/main.py @@ -20,7 +20,7 @@ from dotenv import load_dotenv from constants.colours import LIGHT_YELLOW -from commands import gemini, skullboard, help_menu +from commands import gemini, skullboard, help_menu, admin_commands from utils import time, spam_detection # Load environment variables from .env file @@ -71,6 +71,9 @@ def __init__(self): api_key=GEMINI_API_KEY, ) + self.admin_commands = admin_commands.AdminCommands(gemini_bot=self.gemini_model) + self.tree.add_command(self.admin_commands, guild=Object(GUILD_ID)) + async def setup_hook(self): # Dynamically load all command groups from the commands directory for _, module_name, _ in pkgutil.iter_modules(["src/commands"]): diff --git a/src/models/databases/admin_settings_db.py b/src/models/databases/admin_settings_db.py new file mode 100644 index 0000000..0bed8e8 --- /dev/null +++ b/src/models/databases/admin_settings_db.py @@ -0,0 +1,56 @@ +import os +import sqlite3 +from pathlib import Path +from models.schema.admin_settings_sql import AdminSettingsSQL + + +class AdminSettingsDB: + def __init__(self, db_path: str = "db/admin_settings.db"): + # Ensure the data directory exists + Path(db_path).parent.mkdir(parents=True, exist_ok=True) + + self.db_path = db_path + self.init_db() + + def get_db_connection(self): + return sqlite3.connect(self.db_path) + + def init_db(self): + """Initialise the database with tables and default values from .env""" + with self.get_db_connection() as conn: + cursor = conn.cursor() + + # Create tables + for statement in AdminSettingsSQL.initialisation_tables: + cursor.execute(statement) + + # Initialise with default values from .env + default_settings = { + "GUILD_ID": os.getenv("GUILD_ID", ""), + "SKULLBOARD_CHANNEL_ID": os.getenv("SKULLBOARD_CHANNEL_ID", ""), + "REQUIRED_REACTIONS": os.getenv("REQUIRED_REACTIONS", "3"), + } + + # Only set defaults if the settings don't exist + for key, value in default_settings.items(): + cursor.execute( + "INSERT OR IGNORE INTO settings (key, value) VALUES (?, ?)", + (key, value), + ) + + conn.commit() + + def get_setting(self, key: str) -> str: + """Get a setting value from the database""" + with self.get_db_connection() as conn: + cursor = conn.cursor() + cursor.execute(AdminSettingsSQL.get_setting, (key,)) + result = cursor.fetchone() + return result[0] if result else None + + def set_setting(self, key: str, value: str): + """Set a setting value in the database""" + with self.get_db_connection() as conn: + cursor = conn.cursor() + cursor.execute(AdminSettingsSQL.set_setting, (key, value)) + conn.commit() diff --git a/src/models/schema/admin_settings_sql.py b/src/models/schema/admin_settings_sql.py new file mode 100644 index 0000000..873bfac --- /dev/null +++ b/src/models/schema/admin_settings_sql.py @@ -0,0 +1,21 @@ +class AdminSettingsSQL: + """Store SQL statements for the admin settings functionalities.""" + + initialisation_tables = [ + """ + CREATE TABLE IF NOT EXISTS settings ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL + ); + """ + ] + + get_setting = """ + SELECT value FROM settings WHERE key = ?; + """ + + set_setting = """ + INSERT INTO settings (key, value) + VALUES (?, ?) + ON CONFLICT(key) DO UPDATE SET value = excluded.value; + """ diff --git a/src/utils/settings.py b/src/utils/settings.py new file mode 100644 index 0000000..d07ca57 --- /dev/null +++ b/src/utils/settings.py @@ -0,0 +1,11 @@ +from models.admin_settings_db import AdminSettingsDB +import os + + +def get_setting_with_fallback(key: str, default: str = None) -> str: + """Get a setting from the database, falling back to env vars if not found""" + db = AdminSettingsDB() + value = db.get_setting(key) + if value is None: + value = os.getenv(key, default) + return value