190 lines
8.1 KiB
Python
190 lines
8.1 KiB
Python
import discord
|
|
from discord import app_commands
|
|
import sqlite3
|
|
import os
|
|
from logger import logger
|
|
from data_management.consent import ConsentManager
|
|
from data_management.data_access import DataAccessManager
|
|
import datetime
|
|
import config
|
|
|
|
|
|
class Birthday:
|
|
def __init__(self, client):
|
|
self.client = client
|
|
self.db_path = "data/EONA.db"
|
|
self.logger = logger
|
|
self.consent_manager = ConsentManager(self.db_path)
|
|
self.data_access_manager = DataAccessManager(self.db_path)
|
|
self.birthday_channel_id = None # Initially unset
|
|
|
|
# Ensure that the database and tables are set up
|
|
self.ensure_db()
|
|
self.setup_commands()
|
|
|
|
def ensure_db(self):
|
|
os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
|
|
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS birthdays (
|
|
user_id INTEGER,
|
|
guild_id INTEGER,
|
|
user_name TEXT,
|
|
birthday TEXT,
|
|
PRIMARY KEY (user_id, guild_id)
|
|
);
|
|
""")
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS settings (
|
|
guild_id INTEGER PRIMARY KEY,
|
|
birthday_channel_id INTEGER
|
|
);
|
|
""")
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def set_birthday_channel(self, guild_id, channel_id):
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("""
|
|
INSERT INTO settings (guild_id, birthday_channel_id)
|
|
VALUES (?, ?)
|
|
ON CONFLICT(guild_id)
|
|
DO UPDATE SET birthday_channel_id = excluded.birthday_channel_id;
|
|
""", (guild_id, channel_id))
|
|
conn.commit()
|
|
conn.close()
|
|
self.birthday_channel_id = channel_id # Store in the instance
|
|
|
|
def get_birthday_channel(self, guild_id):
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT birthday_channel_id FROM settings WHERE guild_id = ?", (guild_id,))
|
|
row = cursor.fetchone()
|
|
conn.close()
|
|
|
|
if row and row[0]:
|
|
return row[0] # Return the set channel ID
|
|
else:
|
|
return config.DEFAULT_CHANNEL_ID # Fall back to default from config
|
|
|
|
async def send_birthday_alert(self, guild, message):
|
|
channel_id = self.get_birthday_channel(guild.id)
|
|
channel = guild.get_channel(channel_id)
|
|
if channel:
|
|
await channel.send(message)
|
|
else:
|
|
self.logger.error(f"Channel with ID {channel_id} not found in guild {guild.name}")
|
|
|
|
async def delete_birthday(self, user_id, guild_id):
|
|
self.data_access_manager.delete_data(user_id, guild_id)
|
|
|
|
async def get_birthday(self, user_id, guild_id):
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT birthday FROM birthdays WHERE user_id = ? AND guild_id = ?", (user_id, guild_id))
|
|
row = cursor.fetchone()
|
|
conn.close()
|
|
return row[0] if row else None
|
|
|
|
@tasks.loop(hours=24)
|
|
async def check_birthdays(self):
|
|
today = datetime.date.today().strftime("%Y-%m-%d")
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT user_id, guild_id FROM birthdays WHERE birthday = ?", (today,))
|
|
rows = cursor.fetchall()
|
|
conn.close()
|
|
|
|
for user_id, guild_id in rows:
|
|
guild = self.client.get_guild(int(guild_id))
|
|
if guild:
|
|
user = guild.get_member(int(user_id))
|
|
if user:
|
|
message = f"Happy birthday {user.mention}! 🎉🎂"
|
|
await self.send_birthday_alert(guild, message)
|
|
|
|
async def setup_hook(self):
|
|
self.check_birthdays.start()
|
|
self.logger.info("Birthday module loaded")
|
|
|
|
def setup(self, tree: discord.app_commands.CommandTree):
|
|
@tree.command(name="set_birthday_channel", description="Set the channel for birthday alerts")
|
|
@app_commands.checks.has_permissions(administrator=True)
|
|
async def set_birthday_channel(interaction: discord.Interaction, channel: discord.TextChannel):
|
|
guild_id = interaction.guild.id
|
|
self.set_birthday_channel(guild_id, channel.id)
|
|
await interaction.response.send_message(f"Birthday alerts will now be sent to {channel.mention}.", ephemeral=True)
|
|
|
|
self.client.tree.add_command(set_birthday_channel)
|
|
@tree.command(name="consent", description="Give consent for data storage.")
|
|
async def consent_command(interaction: discord.Interaction):
|
|
self.consent_manager.give_consent(interaction.user.id, interaction.guild.id)
|
|
await interaction.response.send_message(f"{interaction.user.mention}, your consent has been recorded.", ephemeral=True)
|
|
|
|
@tree.command(name="view_data", description="View your stored data.")
|
|
async def view_data_command(interaction: discord.Interaction):
|
|
data = self.data_access_manager.view_user_data(interaction.user.id, interaction.guild.id)
|
|
if data:
|
|
await interaction.response.send_message(f"{interaction.user.mention}, your stored data: {data}", ephemeral=True)
|
|
else:
|
|
await interaction.response.send_message(f"{interaction.user.mention}, no data found.", ephemeral=True)
|
|
|
|
@tree.command(name="delete_data", description="Delete your stored data.")
|
|
async def delete_data_command(interaction: discord.Interaction):
|
|
self.data_access_manager.delete_user_data(interaction.user.id, interaction.guild.id)
|
|
await interaction.response.send_message(f"{interaction.user.mention}, your data has been deleted.", ephemeral=True)
|
|
|
|
@tree.command(name="set_birthday", description="Set your birthday (YYYY-MM-DD)")
|
|
async def set_birthday_command(interaction: discord.Interaction):
|
|
await interaction.response.send_modal(BirthdayModal(self))
|
|
|
|
@tree.command(name="get_birthday", description="Check your birthday")
|
|
async def get_birthday_command(interaction: discord.Interaction):
|
|
user_id = str(interaction.user.id)
|
|
guild_id = str(interaction.guild.id)
|
|
birthday = await self.get_birthday(user_id, guild_id)
|
|
if birthday:
|
|
await interaction.response.send_message(f"Your birthday has been set to {birthday}", ephemeral=True)
|
|
else:
|
|
await interaction.response.send_message("Your birthday has not been set. Please use `/set_birthday` to set your birthday.", ephemeral=True)
|
|
|
|
if not tree.get_command("consent"):
|
|
tree.add_command(consent_command)
|
|
if not tree.get_command("view_data"):
|
|
tree.add_command(view_data_command)
|
|
if not tree.get_command("delete_data"):
|
|
tree.add_command(delete_data_command)
|
|
if not tree.get_command("set_birthday"):
|
|
tree.add_command(set_birthday_command)
|
|
if not tree.get_command("get_birthday"):
|
|
tree.add_command(get_birthday_command)
|
|
|
|
|
|
class BirthdayModal(discord.ui.Modal, title="Set your birthday"):
|
|
birthday = discord.ui.TextInput(label="Your birthday", placeholder="YYYY-MM-DD", style=discord.TextStyle.short, required=True)
|
|
|
|
def __init__(self, birthday_module):
|
|
super().__init__()
|
|
self.birthday_module = birthday_module
|
|
|
|
async def on_submit(self, interaction: discord.Interaction):
|
|
user_id = str(interaction.user.id)
|
|
guild_id = str(interaction.guild.id)
|
|
try:
|
|
# Validate the date format
|
|
birthday_str = self.birthday.value
|
|
birthday = datetime.strptime(birthday_str, "%Y-%m-%d").date()
|
|
|
|
# Proceed to set the birthday if valid
|
|
await self.birthday_module.set_birthday(interaction, user_id, guild_id, birthday_str)
|
|
except ValueError:
|
|
await interaction.response.send_message("Invalid date format. Please enter the date in YYYY-MM-DD format.", ephemeral=True)
|
|
|
|
|
|
def setup(client):
|
|
birthday_module = Birthday(client)
|
|
client.add_cog(birthday_module)
|
|
client.birthday_module = birthday_module |