From 07aabac1e4447c9fdb654999237af92d789a924d Mon Sep 17 00:00:00 2001 From: itskarudo Date: Fri, 29 Nov 2024 16:05:10 +0100 Subject: [PATCH 1/2] feat: Add an API for defining cronjobs This defines an easier API for creating cronjobs that run within the bot. It also adds /cron commands to control individual or all crons fix: please the formatter gods? fix: isort things --- eruditus/app_commands/cron/__init__.py | 86 ++++++++++++++++++++++++++ eruditus/cronjobs/__init__.py | 1 + eruditus/cronjobs/jobs/.gitkeep | 0 eruditus/eruditus.py | 2 + eruditus/lib/types.py | 35 +++++++++++ 5 files changed, 124 insertions(+) create mode 100644 eruditus/app_commands/cron/__init__.py create mode 100644 eruditus/cronjobs/__init__.py create mode 100644 eruditus/cronjobs/jobs/.gitkeep diff --git a/eruditus/app_commands/cron/__init__.py b/eruditus/app_commands/cron/__init__.py new file mode 100644 index 0000000..99c3077 --- /dev/null +++ b/eruditus/app_commands/cron/__init__.py @@ -0,0 +1,86 @@ +from typing import Optional + +import discord +from discord import app_commands +from discord.app_commands import Choice + +from cronjobs import CRONJOBS + + +class Cron(app_commands.Group): + """Manage cron jobs.""" + + def __init__(self, client: discord.Client) -> None: + super().__init__(name="cron") + + self.jobs = {} + + for name, job in CRONJOBS.items(): + job.bind_client(client) + self.jobs[name] = job.create_task() + + async def _job_autocompletion_func( + self, _: discord.Interaction, current: str + ) -> list[Choice[str]]: + suggestions = [] + for job in self.jobs: + if current.lower() in job.lower(): + suggestions.append(Choice(name=job, value=job)) + if len(suggestions) == 25: + break + + return suggestions + + @app_commands.checks.bot_has_permissions(manage_channels=True, manage_roles=True) + @app_commands.checks.has_permissions(manage_channels=True, manage_roles=True) + @app_commands.command() + @app_commands.autocomplete(job=_job_autocompletion_func) # type: ignore + async def start(self, interaction: discord.Interaction, job: Optional[str] = None): + """Start a cron job.""" + + if job is None: + for task in self.jobs.values(): + if not task.is_running(): + task.start() + return await interaction.response.send_message( + "✅ All jobs has been started." + ) + + task = self.jobs.get(job) + if task is None: + return await interaction.response.send_message("No such job.") + + if task.is_running(): + return await interaction.response.send_message( + "This job is already running." + ) + + task.start() + await interaction.response.send_message(f"✅ Job `{job}` has been started.") + + @app_commands.checks.bot_has_permissions(manage_channels=True, manage_roles=True) + @app_commands.checks.has_permissions(manage_channels=True, manage_roles=True) + @app_commands.command() + @app_commands.autocomplete(job=_job_autocompletion_func) # type: ignore + async def stop(self, interaction: discord.Interaction, job: Optional[str] = None): + """Stop a cron job.""" + + if job is None: + for task in self.jobs.values(): + if task.is_running(): + task.cancel() + return await interaction.response.send_message( + "✅ All jobs has been stopped." + ) + + task = self.jobs.get(job) + if task is None: + return await interaction.response.send_message("No such job.") + + if not task.is_running(): + return await interaction.response.send_message( + "This job is already stopped." + ) + + task.cancel() + await interaction.response.send_message(f"✅ Job `{job}` has been stopped.") diff --git a/eruditus/cronjobs/__init__.py b/eruditus/cronjobs/__init__.py new file mode 100644 index 0000000..3da3224 --- /dev/null +++ b/eruditus/cronjobs/__init__.py @@ -0,0 +1 @@ +CRONJOBS = {} diff --git a/eruditus/cronjobs/jobs/.gitkeep b/eruditus/cronjobs/jobs/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/eruditus/eruditus.py b/eruditus/eruditus.py index bb35dc7..ec523b1 100644 --- a/eruditus/eruditus.py +++ b/eruditus/eruditus.py @@ -15,6 +15,7 @@ import config from app_commands.bookmark import Bookmark from app_commands.cipher import Cipher +from app_commands.cron import Cron from app_commands.ctf import CTF from app_commands.ctftime import CTFTime from app_commands.encoding import Encoding @@ -191,6 +192,7 @@ async def setup_hook(self) -> None: self.tree.add_command(TakeNote(), guild=discord.Object(GUILD_ID)) self.tree.add_command(CTF(), guild=discord.Object(GUILD_ID)) self.tree.add_command(Intro(), guild=discord.Object(GUILD_ID)) + self.tree.add_command(Cron(self), guild=discord.Object(GUILD_ID)) # Restore `workon` buttons. for challenge in MONGO[DBNAME][CHALLENGE_COLLECTION].find({"solved": False}): diff --git a/eruditus/lib/types.py b/eruditus/lib/types.py index 71b11ef..a48adbf 100644 --- a/eruditus/lib/types.py +++ b/eruditus/lib/types.py @@ -1,4 +1,12 @@ +import datetime +from abc import ABC, abstractmethod +from dataclasses import dataclass from enum import Enum +from typing import Optional, Sequence, Union + +import discord +from discord.ext import tasks +from discord.utils import MISSING class CPUArchitecture(Enum): @@ -33,3 +41,30 @@ class OSType(Enum): class Privacy(Enum): public = 0 private = 1 + + +@dataclass +class CronJob(ABC): + client: Optional[discord.Client] = None + seconds: float = MISSING + minutes: float = MISSING + hours: float = MISSING + time: Union[datetime.time, Sequence[datetime.time]] = MISSING + count: Optional[int] = None + reconnect: bool = True + + @abstractmethod + async def run(self): + pass + + def bind_client(self, client: discord.Client): + self.client = client + + def create_task(self): + return tasks.loop( + seconds=self.seconds, + minutes=self.minutes, + hours=self.hours, + count=self.count, + reconnect=self.reconnect, + )(self.run) From 88d7f2923c5f600658844574b4c5698946f0ddc7 Mon Sep 17 00:00:00 2001 From: Khaled Touahri <41697719+itskarudo@users.noreply.github.com> Date: Fri, 29 Nov 2024 22:40:50 +0100 Subject: [PATCH 2/2] fix typos in job status messages Co-authored-by: hfz <32499116+hfz1337@users.noreply.github.com> --- eruditus/app_commands/cron/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/eruditus/app_commands/cron/__init__.py b/eruditus/app_commands/cron/__init__.py index 99c3077..fa8358b 100644 --- a/eruditus/app_commands/cron/__init__.py +++ b/eruditus/app_commands/cron/__init__.py @@ -43,7 +43,7 @@ async def start(self, interaction: discord.Interaction, job: Optional[str] = Non if not task.is_running(): task.start() return await interaction.response.send_message( - "✅ All jobs has been started." + "✅ All jobs have been started." ) task = self.jobs.get(job) @@ -70,7 +70,7 @@ async def stop(self, interaction: discord.Interaction, job: Optional[str] = None if task.is_running(): task.cancel() return await interaction.response.send_message( - "✅ All jobs has been stopped." + "✅ All jobs have been stopped." ) task = self.jobs.get(job)