nixos-config/sites/sunflower/reminder_bot.py

174 lines
6.6 KiB
Python

from cronsim import CronSimError
import discord
from discord.ext import commands
import aiosqlite
import aiocron
from datetime import datetime, timedelta, timezone
import logging
log = logging.getLogger("reminder_bot")
DATABASE = '/var/lib/reminder-bot/db'
with open('/var/lib/reminder-bot/token') as fp:
TOKEN = fp.read().strip()
BOTSPAM_CHANNEL_ID = 1428820508178124820
intents = discord.Intents.default()
intents.message_content = True
client = commands.Bot('/', intents=intents)
tree = client.tree
# ------------------------------------------------------------
# DATABASE SETUP
# ------------------------------------------------------------
async def init_db():
async with aiosqlite.connect(DATABASE) as db:
await db.execute("""
CREATE TABLE IF NOT EXISTS reminders (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message TEXT NOT NULL,
cron TEXT NOT NULL
)
""")
await db.commit()
# ------------------------------------------------------------
# DISMISS BUTTON
# ------------------------------------------------------------
class DismissButton(discord.ui.View):
def __init__(self, reminder_id: int):
super().__init__(timeout=None)
self.reminder_id = reminder_id
@discord.ui.button(label="Dismiss", style=discord.ButtonStyle.red)
async def dismiss(self, interaction: discord.Interaction, button: discord.ui.Button):
user = interaction.user.display_name
timestamp = datetime.now(tz=timezone(timedelta(hours=-7))).strftime("%a %b %-m %-I:%M %p")
assert interaction.message is not None
button.disabled = True
await interaction.response.edit_message(content=interaction.message.content + f"\nCompleted by {user} at {timestamp}", view=self)
# ------------------------------------------------------------
# GLOBAL JOB REGISTRY
# ------------------------------------------------------------
scheduled_jobs = {}
async def send_reminder(message: str, reminder_id: int):
channel = client.get_channel(BOTSPAM_CHANNEL_ID)
assert isinstance(channel, discord.TextChannel)
await channel.send(
f"⏰ **Reminder:** {message}",
view=DismissButton(reminder_id),
allowed_mentions=discord.AllowedMentions(everyone=True, users=True, roles=True)
)
async def schedule_reminder(reminder_id: int, message: str, cron_expr: str):
assert reminder_id not in scheduled_jobs
job = aiocron.crontab(cron_expr, func=send_reminder, args=(message, reminder_id))
scheduled_jobs[reminder_id] = job
# ------------------------------------------------------------
# COMMANDS
# ------------------------------------------------------------
@tree.command(name="reminder-add")
async def add_reminder(ctx: discord.Interaction, cron_expr: str, *, message: str):
"""
Add a new reminder using a cron expression.
"""
try:
aiocron.CronSim(cron_expr, datetime.now())
except CronSimError as e:
await ctx.response.send_message(f"Your cron expression did not parse: {e}", ephemeral=True)
else:
async with aiosqlite.connect(DATABASE) as db:
cursor = await db.execute(
"INSERT INTO reminders (message, cron) VALUES (?, ?)",
(message, cron_expr)
)
await db.commit()
reminder_id = cursor.lastrowid
assert reminder_id is not None
await schedule_reminder(reminder_id, message, cron_expr)
await ctx.response.send_message(f"✅ Reminder {reminder_id} added: '{message}' `{cron_expr}`")
@tree.command(name="reminder-list")
async def list_reminders(ctx: discord.Interaction):
"""List all active reminders."""
async with aiosqlite.connect(DATABASE) as db:
async with db.execute("SELECT id, message, cron FROM reminders") as cursor:
rows = await cursor.fetchall()
if not rows:
await ctx.response.send_message("No active reminders.", ephemeral=True)
return
lines = [f"**{r[0]}**: {r[1]} `{r[2]}`" for r in rows]
await ctx.response.send_message("\n".join(lines), ephemeral=True)
@tree.command(name="reminder-delete")
async def delete_reminder(ctx: discord.Interaction, reminder_id: int):
"""Delete a reminder by ID."""
async with aiosqlite.connect(DATABASE) as db:
async with db.execute("SELECT message FROM reminders WHERE id = ?", (reminder_id,)) as cursor:
rows = list(await cursor.fetchall())
if not rows:
await ctx.response.send_message(f"No such reminder {reminder_id}")
return
await db.execute("DELETE FROM reminders WHERE id = ?", (reminder_id,))
await db.commit()
job = scheduled_jobs.pop(reminder_id, None)
if job:
job.stop()
await ctx.response.send_message(f"🗑️ Reminder {reminder_id} ({rows[0][0]}) deleted.")
@tree.command(name="reminder-edit")
async def edit_reminder(ctx: discord.Interaction, reminder_id: int, message: str, cron_expr: str):
try:
aiocron.CronSim(cron_expr, datetime.now())
except CronSimError as e:
await ctx.response.send_message(f"Your cron expression did not parse: {e}", ephemeral=True)
else:
async with aiosqlite.connect(DATABASE) as db:
async with db.execute("SELECT message FROM reminders WHERE id = ?", (reminder_id,)) as cursor:
rows = list(await cursor.fetchall())
if not rows:
await ctx.response.send_message(f"No such reminder {reminder_id}")
return
await db.execute("UPDATE reminders SET message = ?, cron = ? WHERE id = ?", (message, cron_expr, reminder_id))
await db.commit()
job = scheduled_jobs.pop(reminder_id, None)
if job:
job.stop()
await schedule_reminder(reminder_id, message, cron_expr)
await ctx.response.send_message(f"✅ Reminder {reminder_id} updated: {message} `{cron_expr}`")
# ------------------------------------------------------------
# STARTUP
# ------------------------------------------------------------
@client.event
async def on_ready():
await init_db()
# Load all reminders from DB and schedule them
async with aiosqlite.connect(DATABASE) as db:
async with db.execute("SELECT id, message, cron FROM reminders") as cursor:
async for reminder_id, message, cron_expr in cursor:
await schedule_reminder(reminder_id, message, cron_expr)
log.info(f"Loaded {len(scheduled_jobs)} scheduled reminders.")
await tree.sync()
# ------------------------------------------------------------
# RUN
# ------------------------------------------------------------
client.run(TOKEN)