Added toggles for spam detection and flood

wbb/modules/filters.py

@@ -60,7 +60,7 @@ async def save_filters(_, message):
     name = message.text.split(None, 1)[1].strip()
     if not name:
         return await message.reply_text(
-            "**Usage**\n__/filter [FILTER_NAME]__"
+            "**Usage:**\n__/filter [FILTER_NAME]__"
         )
     chat_id = message.chat.id
     _type = "text" if message.reply_to_message.text else "sticker"
@@ -96,12 +96,12 @@ async def get_filterss(_, message):
 async def del_filter(_, message):
     if len(message.command) < 2:
         return await message.reply_text(
-            "**Usage**\n__/stop [FILTER_NAME]__"
+            "**Usage:**\n__/stop [FILTER_NAME]__"
         )
     name = message.text.split(None, 1)[1].strip()
     if not name:
         return await message.reply_text(
-            "**Usage**\n__/stop [FILTER_NAME]__"
+            "**Usage:**\n__/stop [FILTER_NAME]__"
         )
     chat_id = message.chat.id
     deleted = await delete_filter(chat_id, name)

wbb/modules/flood.py

@@ -31,18 +31,20 @@ from pyrogram.types import (CallbackQuery, ChatPermissions,
 
 from wbb import SUDOERS, app
 from wbb.core.decorators.errors import capture_err
+from wbb.core.decorators.permissions import adminsOnly
 from wbb.modules.admin import list_admins, member_permissions
+from wbb.utils.dbfunctions import flood_off, flood_on, is_flood_on
 from wbb.utils.filter_groups import flood_group
 
 __MODULE__ = "Flood"
 __HELP__ = """
 Anti-Flood system, the one who sends more than 10 messages in a row, gets muted for an hour (Except for admins).
 
-And no, you can't change the number of messages or action type.
+/flood [ENABLE|DISABLE] - Turn flood detection on or off
 """
 
 
-DB = {}  # NOTE Use mongodb instead of a fucking dict.
+DB = {}  # TODO Use mongodb instead of a fucking dict.
 
 
 def reset_flood(chat_id, user_id=0):
@@ -63,7 +65,8 @@ def reset_flood(chat_id, user_id=0):
 @capture_err
 async def flood_control_func(_, message: Message):
     chat_id = message.chat.id
-
+    if not (await is_flood_on(chat_id)):
+        return
     # Initialize db if not already.
     if chat_id not in DB:
         DB[chat_id] = {}
@@ -143,3 +146,25 @@ async def flood_callback_func(_, cq: CallbackQuery):
     text = f"~~{text}~~\n\n"
     text += f"__User unmuted by {from_user.mention}__"
     await cq.message.edit(text)
+
+
+@app.on_message(filters.command("flood") & ~filters.private)
+@adminsOnly("can_change_info")
+async def flood_toggle(_, message: Message):
+    if len(message.command) != 2:
+        return await message.reply_text(
+            "Usage: /flood [ENABLE|DISABLE]"
+        )
+    status = message.text.split(None, 1)[1].strip()
+    status = status.lower()
+    chat_id = message.chat.id
+    if status == "enable":
+        await flood_on(chat_id)
+        await message.reply_text("Enabled Flood Checker.")
+    elif status == "disable":
+        await flood_off(chat_id)
+        await message.reply_text("Disabled Flood Checker.")
+    else:
+        await message.reply_text(
+            "Unknown Suffix, Use /flood [ENABLE|DISABLE]"
+        )

wbb/modules/spam.py

@@ -26,8 +26,12 @@ from pyrogram.types import (CallbackQuery, InlineKeyboardButton,
                             InlineKeyboardMarkup, Message)
 
 from wbb import LOG_GROUP_ID, SUDOERS, app
+from wbb.core.decorators.permissions import adminsOnly
 from wbb.modules.admin import list_admins, member_permissions
 from wbb.modules.trust import get_spam_data
+from wbb.utils.dbfunctions import (is_spam_detection_on,
+                                   spam_detection_off,
+                                   spam_detection_on)
 from wbb.utils.filter_groups import spam_protection_group
 
 __MODULE__ = "AntiSpam"
@@ -44,7 +48,10 @@ __HELP__ = """
         To mark a message as spam, this will help devs to
         improve spam protection algorithm.
 
-As of now, you cannot turn this off, but we'll add an enable/disable command in the future
+    - /spam_detection [ENABLE|DISABLE]
+        To turn off spam detection in your chat.
+
+**Plot Twist:** NO, WE'RE NOT USING INTELLIVOID'S SPAM DETECTION API!
 """
 
 
@@ -62,6 +69,9 @@ async def spam_protection_func(_, message: Message):
     # We'll handle admins only if it's spam, ignore only sudo users for now.
     if user.id in SUDOERS:
         return
+    enabled = await is_spam_detection_on(chat_id)
+    if not enabled:
+        return
     data = await get_spam_data(message, text)
     if isinstance(data, str):
         return
@@ -70,7 +80,7 @@ async def spam_protection_func(_, message: Message):
     if user.id in (await list_admins(chat_id)):
         return
     text = f"""
-🚨  **SPAM DETECTED** 🚨
+🚨  **SPAM DETECTED**  🚨
 
 **User:** {user.mention}
 **Message:** [Link]({message.link})
@@ -166,7 +176,7 @@ async def spam_flag_func(_, message: Message):
         )
 
     msg = f"""
-**ADMINS OF {message.chat.id} FLAGGED THIS MESSAGE AS SPAM**
+**ADMINS OF {message.chat.id} FLAGGED THIS MESSAGE AS SPAM. [Suggestion]**
 
 ```
 {text.markdown}
@@ -180,3 +190,25 @@ __{dev_forward}__
     await app.send_message(
         LOG_GROUP_ID, msg, disable_web_page_preview=True
     )
+
+
+@app.on_message(filters.command("spam_detection") & ~filters.private)
+@adminsOnly("can_change_info")
+async def spam_toggle(_, message: Message):
+    if len(message.command) != 2:
+        return await message.reply_text(
+            "Usage: /spam_detection [ENABLE|DISABLE]"
+        )
+    status = message.text.split(None, 1)[1].strip()
+    status = status.lower()
+    chat_id = message.chat.id
+    if status == "enable":
+        await spam_detection_on(chat_id)
+        await message.reply_text("Enabled Spam Detection System.")
+    elif status == "disable":
+        await spam_detection_off(chat_id)
+        await message.reply_text("Disabled Spam Detection System.")
+    else:
+        await message.reply_text(
+            "Unknown Suffix, Use /spam_detection [ENABLE|DISABLE]"
+        )

wbb/modules/tts.py

@@ -32,9 +32,6 @@ from pyrogram.types import Message
 
 from wbb import app
 
-__MODULE__ = "TTS"
-__HELP__ = "/tts - Convert Text To Speech."
-
 
 def convert(text):
     audio = BytesIO()

wbb/utils/dbfunctions.py

@@ -26,6 +26,13 @@ from typing import Dict, List, Union
 from wbb import db
 from wbb.utils.functions import obj_to_str, str_to_obj
 
+"""
+SOME THINGS ARE FUCKED UP HERE, LIKE TOGGLEABLES HAVE THEIR OWN COLLECTION
+(SHOULD FIX IT WITH SOMETHING LIKE TOGGLEDB, BUT I WON'T, AS IT WILL TAKE 
+TOO MUCH TIME AND WILL BE BAD FOR ALREADY STORED DATA)
+"""
+
+
 notesdb = db.notes
 filtersdb = db.filters
 warnsdb = db.warns
@@ -46,6 +53,8 @@ sudoersdb = db.sudoers
 blacklist_chatdb = db.blacklistChat
 restart_stagedb = db.restart_stage
 trustdb = db.trust
+flood_toggle_db = db.flood_toggle
+spam_toggle_db = db.spam_toggle
 
 """ Notes functions """
 
@@ -788,3 +797,51 @@ async def update_trust_db(user_id: int, new_data: float):
         {"$set": {"data": data}},
         upsert=True,
     )
+
+
+"""SPAM DETECTION System"""
+
+
+async def is_spam_detection_on(chat_id: int) -> bool:
+    chat = await spam_toggle_db.find_one({"chat_id": chat_id})
+    if not chat:
+        return True
+    return False
+
+
+async def spam_detection_on(chat_id: int):
+    is_on = await is_spam_detection_on(chat_id)
+    if is_on:
+        return
+    return await spam_toggle_db.delete_one({"chat_id": chat_id})
+
+
+async def spam_detection_off(chat_id: int):
+    is_on = await is_spam_detection_on(chat_id)
+    if not is_on:
+        return
+    return await spam_toggle_db.insert_one({"chat_id": chat_id})
+
+
+"""FLOOD System"""
+
+
+async def is_flood_on(chat_id: int) -> bool:
+    chat = await flood_toggle_db.find_one({"chat_id": chat_id})
+    if not chat:
+        return True
+    return False
+
+
+async def flood_on(chat_id: int):
+    is_flood = await is_flood_on(chat_id)
+    if is_flood:
+        return
+    return await flood_toggle_db.delete_one({"chat_id": chat_id})
+
+
+async def flood_off(chat_id: int):
+    is_flood = await is_flood_on(chat_id)
+    if not is_flood:
+        return
+    return await flood_toggle_db.insert_one({"chat_id": chat_id})