From ddc6068b7a01cbed51e2e41a22c6cc7a23e83fff Mon Sep 17 00:00:00 2001
From: Sweetbread <glebzaharov1@gmail.com>
Date: Thu, 18 May 2023 22:55:17 +0300
Subject: [PATCH] [feat]: Add bayan checker

---
 cogs/settings/server.py | 77 ++++++++++++++++++++++++++++++++++++-----
 1 file changed, 69 insertions(+), 8 deletions(-)

diff --git a/cogs/settings/server.py b/cogs/settings/server.py
index 3188353..cca16a0 100644
--- a/cogs/settings/server.py
+++ b/cogs/settings/server.py
@@ -1,3 +1,4 @@
+import discord
 from discord import TextChannel, app_commands, Interaction
 from discord.ext import commands, tasks
 
@@ -7,15 +8,25 @@ import schedule
 from loguru import logger
 from datetime import datetime
 
+import requests
+from io import BytesIO
+from hashlib import sha256
+from os.path import splitext
+
 col = db.guild_settings
 
 @app_commands.guild_only()
 @app_commands.default_permissions(administrator=True)
 class ServerSettings(app_commands.Group, name="настройки_сервера"):
-    def add_server_settings(self, server_id: int):
-        if not col.find_one({"id": server_id}):
-            col.insert_one({
+    def __init__(self, bot):
+        super().__init__()
+        bot.add_listener(self.bayans_listener, name='on_message')
+
+    async def add_server_settings(self, server_id: int):
+        if not (await col.find_one({"id": server_id})):
+            await col.insert_one({
                 "id": server_id,
+                'type': 'general',
                 "levelup": "none",
                 "flood_channels": [],
                 "commands_channels": [],
@@ -42,12 +53,12 @@ class ServerSettings(app_commands.Group, name="настройки_сервера
 
     @app_commands.command(name="назначить_флудилкой")
     async def add_flood_channel(self, inter: Interaction, channel: TextChannel = None):
-        self.add_server_settings(inter.guild_id)
+        await self.add_server_settings(inter.guild_id)
 
         if channel is None:
             channel = inter.channel
         
-        col.update_one(
+        await col.update_one(
             {"id": inter.guild_id},
             {"$push": {"flood_channels": channel.id}}
         )
@@ -56,14 +67,64 @@ class ServerSettings(app_commands.Group, name="настройки_сервера
 
     @app_commands.command(name="управление_ответами")
     async def set_aux_talk(self, inter: Interaction, status: bool):
-        self.add_server_settings(inter.guild_id)
+        await self.add_server_settings(inter.guild_id)
 
-        col.update_one({'id': inter.guild_id}, {'$set': {'aux.talk': status}})
+        await col.update_one({'id': inter.guild_id}, {'$set': {'aux.talk': status}})
         if status:
             await inter.response.send_message("Теперь бот будет отвечать")
         else:
             await inter.response.send_message("Бот больше не будет отвечать")
 
+    @app_commands.command()
+    async def check_bayans(self, itr: Interaction, channel: TextChannel = None):
+        if not channel:
+            channel = itr.channel
+
+        search_dict = {'id': itr.guild_id, 'type': 'bayans', 'channel': channel.id}
+
+        if (await col.find_one(search_dict)):
+            await col.delete_one(search_dict)
+            await db.bayans.delete_many({'channel': channel.id})
+            await itr.response.send_message(f"В {channel.mention} больше не стоит проверка на баяны", delete_after=15)
+        else:
+            await col.insert_one(search_dict)
+            string = "Обработка всех сообщений... Это может занять несколько лет\n{:>4} сообщений обработано"
+            await itr.response.send_message("Обработка всех сообщений... Это может занять несколько лет")
+
+            i = 0
+            async for mes in channel.history(limit=None):
+                if i%1000 == 0:
+                    await itr.edit_original_response(content=string.format(i))
+
+                for attachment in filter(lambda m: splitext(m.filename)[1] in ('.png', '.jpg'), mes.attachments):
+                    img = requests.get(attachment.url)
+                    with BytesIO(img.content) as buffer:
+                        hash0 = sha256(buffer.getbuffer()).hexdigest()
+                    if not (await db.bayans.find_one({'id': mes.guild.id, 'channel': mes.channel.id, 'sha256': hash0})):
+                        await db.bayans.insert_one({'id': mes.guild.id, 'channel': mes.channel.id, 'message': mes.id, 'sha256': hash0, 'index': i})
+
+                i += 1
+
+            await itr.edit_original_response(content=f"Теперь в {channel.mention} стоит проверка на баяны. Сообщений обработано: {i}", delete_after=15)
+
+    async def bayans_listener(self, message: discord.Message):
+        if message.attachments and (await col.find_one({'id': message.guild.id, 'type': 'bayans', 'channel': message.channel.id})):
+            out = []
+
+            for i, attachment in enumerate(filter(lambda m: splitext(m.filename)[1] in ('.png', '.jpg'), message.attachments)):
+                img = requests.get(attachment.url)
+                with BytesIO(img.content) as buffer:
+                    hash0 = sha256(buffer.getbuffer()).hexdigest()
+                if (a := await db.bayans.find_one({'id': message.guild.id, 'channel': message.channel.id, 'sha256': hash0})):
+                    out.append(f"Найдено {i+1} изображение в https://discord.com/channels/{a['id']}/{a['channel']}/{a['message']} /{a['index']+1} (SHA256: `{hash0}`)")
+                else:
+                    await db.bayans.insert_one({'id': message.guild.id, 'channel': message.channel.id, 'message': message.id, 'sha256': hash0, 'index': i})
+            
+            if out:
+                await message.reply('\n'.join(out))
+
+
+
 #class NiceName(commands.Cog):
 #    def __init__(self, bot):
 #        self.bot = bot
@@ -94,4 +155,4 @@ class ServerSettings(app_commands.Group, name="настройки_сервера
 
 
 async def setup(bot):
-    bot.tree.add_command(ServerSettings())
+    bot.tree.add_command(ServerSettings(bot))