diff --git a/jarvis_tasks/__init__.py b/jarvis_tasks/__init__.py index 64e2bc7..2d0adef 100644 --- a/jarvis_tasks/__init__.py +++ b/jarvis_tasks/__init__.py @@ -16,6 +16,7 @@ from jarvis_tasks.tasks import ( lockdown, reddit, reminder, + temprole, twitter, warning, ) @@ -62,6 +63,7 @@ async def _start(config: Optional[str] = "config.yaml") -> None: lockdown.lift, reddit.reddit, reminder.remind, + temprole.remove, twitter.twitter, warning.unwarn, ] diff --git a/jarvis_tasks/tasks/temprole.py b/jarvis_tasks/tasks/temprole.py new file mode 100644 index 0000000..27815ec --- /dev/null +++ b/jarvis_tasks/tasks/temprole.py @@ -0,0 +1,66 @@ +"""JARVIS warnings tasks.""" +import asyncio +import logging +from datetime import datetime, timedelta, timezone + +from dis_snek import Snake +from jarvis_core.db import q +from jarvis_core.db.models import Temprole + +from jarvis_tasks.util import runat + +queue = [] +logger = logging.getLogger(__name__) + + +async def _remove(bot: Snake, temprole: Temprole) -> None: + logger.debug(f"Removing temprole {temprole.id}") + guild = await bot.fetch_guild(temprole.guild) + if not guild: + logger.warn(f"Temprole {temprole.id} invalid") + await temprole.delete() + return + + user = await guild.fetch_member(temprole.user) + if not user: + logger.warn(f"Temprole {temprole.id} invalid") + await temprole.delete() + return + + role = await guild.fetch_role(temprole.role) + if not role: + logger.warn(f"Temprole {temprole.id} invalid") + await temprole.delete() + return + + if user.has_role(role): + try: + await user.remove_role(role, "Temporary role expired") + except Exception: + logger.warn("Temprole failed to remove, ignoring") + + await temprole.delete() + + +async def remove(bot: Snake) -> None: + """ + Remove temproles when they expire. + + Args: + bot: Snake instance + """ + logger.debug("Starting Task-remove") + while True: + max_ts = datetime.now(tz=timezone.utc) + timedelta(minutes=55) + temproles = Temprole.find(q(expires_at__lte=max_ts, id__nin=queue)) + async for temprole in temproles: + if temprole.id in queue: + logger.warn("Temprole found despite filter") + continue + + coro = _remove(bot, temprole) + when = temprole.expires_at + asyncio.create_task(runat(when, coro, logger)) + queue.append(temprole.id) + # Check every hour + await asyncio.sleep(3600) diff --git a/pyproject.toml b/pyproject.toml index 3d711c6..25176c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "jarvis-tasks" -version = "0.6.1" +version = "0.7.0" description = "" authors = ["Your Name "]