jarvis-tasks/jarvis_tasks/tasks/ban.py
2023-05-10 18:06:44 -06:00

63 lines
1.9 KiB
Python

"""JARVIS ban tasks."""
import asyncio
import logging
from datetime import datetime, timedelta, timezone
from beanie.operators import LTE
from jarvis_core.db.models import Ban, Unban
from interactions import Client
from interactions.client.errors import NotFound
from interactions.models.discord.guild import Guild
from interactions.models.discord.user import User
from jarvis_tasks.util import runat
queue = []
logger = logging.getLogger(__name__)
async def _unban(bot: int, guild: Guild, user: User, ban: Ban) -> None:
if guild and user:
logger.debug(f"Unbanning user {user.id} from guild {guild.id}")
try:
await guild.unban(user=user, reason="JARVIS tempban expired")
except NotFound:
logger.debug(f"User {user.id} not banned from guild {guild.id}")
ban.active = False
await ban.save()
await Unban(
user=user.id,
guild=guild.id,
username=user.username,
discrim=user.discriminator,
admin=bot,
reason="Ban expired",
).save()
queue.remove(ban.id)
async def unban(bot: Client) -> None:
"""
Unban users when ban expires.
Args:
bot: Client instance
"""
logger.debug("Starting Task-ban")
while True:
max_ts = datetime.now(tz=timezone.utc) + timedelta(minutes=9)
bans = Ban.find(
Ban.type == "temp", Ban.active == True, LTE(Ban.duration, max_ts)
)
async for ban in bans:
if ban.id in queue:
continue
guild = await bot.fetch_guild(ban.guild)
user = await bot.fetch_user(ban.user)
coro = _unban(bot.user.id, guild, user, ban, logger)
when = ban.created_at + timedelta(hours=ban.duration)
asyncio.create_task(runat(when, coro, logger))
queue.append(ban.id)
# Check ever 10 minutes
await asyncio.sleep(600)