discord-cleanup-bot/bot.py

160 lines
7.3 KiB
Python
Raw Normal View History

2023-06-25 00:50:23 +02:00
import os, logging
from datetime import datetime, timedelta
from typing import Optional
from sqlalchemy import create_engine
from sqlalchemy.orm import Session
from models import Guild, Channel, Base
import discord
2023-10-28 01:44:20 +02:00
from discord import app_commands, Game, Permissions
2023-06-25 00:50:23 +02:00
from discord.app_commands import Choice
from discord.ext import tasks
2023-10-28 01:44:20 +02:00
from discord.utils import oauth_url
2023-06-25 00:50:23 +02:00
logger = logging.getLogger()
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
dt_fmt = '%Y-%m-%d %H:%M:%S'
formatter = logging.Formatter('[{asctime}] [{levelname:<8}] {name}: {message}', dt_fmt, style='{')
handler.setFormatter(formatter)
logger.addHandler(handler)
sqla_engine = create_engine(f'sqlite:///{os.getenv("DB_PATH", "bot.db")}')
Base.metadata.create_all(sqla_engine)
# Yield successive n-sized
# chunks from l.
def divide_chunks(l, n):
# looping till length l
for i in range(0, len(l), n):
yield l[i:i + n]
2023-06-25 23:54:07 +02:00
async def purge_messages(channel, audit_message: str, before: datetime, after: datetime):
messages = [message async for message in channel.history(before=before, after=after)]
filtered_messages = [m for m in messages if m.pinned == False]
message_chunks = divide_chunks(filtered_messages, 100)
for chunk in message_chunks:
await channel.delete_messages(chunk, reason=audit_message)
2023-06-25 00:50:23 +02:00
class BotClient(discord.Client):
def __init__(self, *, intents: discord.Intents):
super().__init__(intents=intents)
self.tree = app_commands.CommandTree(self)
async def setup_hook(self):
await self.tree.sync()
self.background_task.start()
2023-10-28 01:44:20 +02:00
invite_url = oauth_url(client_id=self.application_id, permissions=Permissions(permissions=76800))
logger.info(f"Bot started. Invite it using the following link: {invite_url}")
2023-06-25 00:50:23 +02:00
@tasks.loop(minutes=15)
async def background_task(self):
with Session(sqla_engine) as session:
persisted_channels = session.query(Channel).order_by(Channel.last_pruned.desc()).limit(5).all()
for persisted_channel in persisted_channels:
channel = self.get_channel(persisted_channel.channel_id)
logger.info(f"Cleaning channel {persisted_channel.channel_id} in guild {persisted_channel.guild_id} ({channel.guild.name})")
now = datetime.now()
before = now - timedelta(hours=persisted_channel.retention_hours)
after = now - timedelta(days=14)
2023-06-25 23:54:07 +02:00
await purge_messages(channel=channel, audit_message='Configured retention period expired.', before=before, after=after)
2023-06-25 00:50:23 +02:00
persisted_channel.last_pruned=now
session.commit()
@background_task.before_loop
async def before_my_task(self):
await self.wait_until_ready()
intents = discord.Intents.default()
client = BotClient(intents=intents)
@client.event
async def on_guild_join(guild):
with Session(sqla_engine) as session:
new_guild = Guild(guild_id=guild.id)
session.add(new_guild)
session.commit()
@client.event
async def on_guild_remove(guild):
with Session(sqla_engine) as session:
removed_guild = session.get(Guild, guild.id)
if removed_guild == None:
return
session.delete(removed_guild)
session.commit()
2023-06-25 01:53:12 +02:00
@client.event
async def on_guild_channel_delete(channel):
with Session(sqla_engine) as session:
removed_channel = session.get(Channel, channel.id)
if removed_channel == None:
return
session.delete(removed_channel)
session.commit()
2023-06-25 00:50:23 +02:00
@client.tree.command()
@app_commands.guild_only()
@app_commands.describe(
action='Action to perform',
retention_period='Use thogether with the "set" action to define a retention period in days',
)
@app_commands.choices(action=[
Choice(name='get', value=1),
Choice(name='set', value=2),
Choice(name='disable', value=3),
])
async def retention(interaction: discord.Interaction, action: Choice[int], retention_period: Optional[int] = None):
"""Manage retention settings for a channel."""
match action.value:
case 1:
with Session(sqla_engine) as session:
persisted_channel = session.get(Channel, interaction.channel_id)
if persisted_channel == None:
await interaction.response.send_message(content=' There is currently no retention period set for this channel.', delete_after=20)
return
await interaction.response.send_message(content=f' Messages in this channel will currently be deleted after `{persisted_channel.retention_hours / 24}` days.', delete_after=20)
case 2:
if retention_period == None:
await interaction.response.send_message(content='❌ Error: You need to specify a `retention_period` when using the `set` action', delete_after=20)
return
if retention_period > 13:
await interaction.response.send_message(content='❌ Error: Due to technical limitations, 13 days is the maximum after which I can still delete messages. Please use `13` or less as the retention period.', delete_after=20)
return
with Session(sqla_engine) as session:
channel = session.get(Channel, interaction.channel_id)
if channel == None:
channel = Channel(channel_id=interaction.channel_id, guild_id=interaction.guild_id, retention_hours=(retention_period * 24), last_pruned=datetime.min)
session.add(channel)
else:
channel.retention_hours = retention_period * 24
session.commit()
await interaction.response.send_message(content=f'✅ Messages in this channel will be automatically deleted after `{retention_period}` days.', delete_after=20)
case 3:
with Session(sqla_engine) as session:
channel = session.get(Channel, interaction.channel_id)
if channel != None:
session.delete(channel)
session.commit()
await interaction.response.send_message(content='✅ Messages in this channel will no loger be automatically deleted.', delete_after=20)
2023-06-25 23:54:07 +02:00
@client.tree.command()
@app_commands.guild_only()
@app_commands.describe(
days='Delete messages of the last x days'
)
async def purge(interaction: discord.Interaction, days: int):
"""Clear messages of the last x days in the current channel."""
if days > 13:
await interaction.response.send_message(content='❌ Error: Due to technical limitations, 13 days is the maximum after which I can still delete messages. Please use `13` days or less.', delete_after=20)
return
2023-07-11 23:13:04 +02:00
logger.info(f"Purging messages of the last {days} days from channel {interaction.channel_id} ({interaction.channel.name}) in guild {interaction.guild_id} ({interaction.guild.name}) at request of {interaction.user.name}.")
2023-07-11 23:07:32 +02:00
await interaction.response.send_message(content='✅ Messages will be deleted shortly.', delete_after=20)
await purge_messages(channel=interaction.channel, audit_message=f'Deleted via /purge command from user {interaction.user.name}', before=interaction.created_at, after=(datetime.now() - timedelta(days=days)))
2023-06-25 23:54:07 +02:00
2023-07-03 22:01:17 +02:00
client.activity = Game(name="with the broomstick")
2023-07-02 20:40:39 +02:00
2023-06-25 00:50:23 +02:00
client.run(token=os.getenv("BOT_TOKEN"), log_handler=None)