jarvis-tasks/jarvis_tasks/tasks/lock.py

54 lines
1.9 KiB
Python

"""JARVIS Lock background task."""
import asyncio
from datetime import datetime, timedelta, timezone
from logging import Logger
from dis_snek import Snake
from dis_snek.client.utils.misc_utils import get
from dis_snek.models.discord.channel import GuildChannel
from jarvis_core.db import q
from jarvis_core.db.models import Lock
from jarvis_tasks.util import runat
async def _unlock(channel: GuildChannel, lock: Lock, logger: Logger) -> None:
logger.debug(f"Deactivating lock {lock.id}")
try:
overwrite = get(channel.permission_overwrites, id=lock.guild)
if overwrite and lock.original_perms:
overwrite.allow = lock.original_perms.allow
overwrite.deny = lock.original_perms.deny
try:
await channel.edit_permission(overwrite, reason="Automatic unlock")
except Exception:
logger.debug("Locked channel deleted, ignoring error")
elif overwrite and not lock.original_perms:
await channel.delete_permission(target=overwrite, reason="Automatic unlock")
else:
logger.debug("Permission neither exists not existed")
lock.active = False
except Exception:
logger.debug("Locked channel deleted, ignoring error")
await lock.commit()
async def unlock(bot: Snake, logger: Logger) -> None:
"""
Unlock locked channels.
Args:
bot: Bot instance
logger: Global logger
"""
while True:
max_ts = datetime.now(tz=timezone.utc) + timedelta(seconds=55)
locks = Lock.find(q(active=True, created_at__lte=max_ts))
async for lock in locks:
guild = await bot.fetch_guild(lock.guild)
channel = await guild.fetch_channel(lock.channel)
coro = _unlock(channel, lock, logger)
when = lock.created_at + timedelta(minutes=lock.duration)
asyncio.create_task(runat(when, coro, logger))
await asyncio.sleep(delay=60)