Some small fixes

wbb/__init__.py

@@ -31,9 +31,10 @@ from os import path
 from aiohttp import ClientSession
 from motor.motor_asyncio import AsyncIOMotorClient as MongoClient
 from pyrogram import Client
-from pyromod import listen
 from Python_ARQ import ARQ
 
+from pyromod import listen
+
 # Setup logging
 log_file = "error.log"
 
@@ -108,7 +109,6 @@ aiohttpsession = ClientSession()
 # ARQ Client
 print("[INFO]: INITIALIZING ARQ CLIENT")
 arq = ARQ(ARQ_API_URL, ARQ_API_KEY, aiohttpsession)
-listen.listen.PyroARQ(arq)
 # Bot client
 print("[INFO]: INITIALIZING BOT CLIENT")
 app = Client("wbb", bot_token=BOT_TOKEN, api_id=API_ID, api_hash=API_HASH)
@@ -158,7 +158,7 @@ def get_info(app, app2):
         for i in app2.iter_dialogs()
         if i.chat.type in ["group", "supergroup", "private"]
     ]
-    ub_b_common_chats = [i.id for i in app2.get_common_chats(BOT_ID)]
+    ub_b_common_chats = [i.id for i in app2.get_common_chats(BOT_USERNAME)]
     USERBOT_BOT_CHAT_DIFFERENCE = list(
         set(all_ub_chats).difference(ub_b_common_chats)
     )

wbb/modules/admin.py

@@ -80,12 +80,12 @@ from wbb.core.decorators.permissions import adminsOnly
 
 
 async def list_admins(chat_id: int):
-    list_of_admins = []
-    async for member in app.iter_chat_members(
-        chat_id, filter="administrators"
-    ):
-        list_of_admins.append(member.user.id)
-    return list_of_admins
+    return [
+        member.user.id
+        async for member in app.iter_chat_members(
+            chat_id, filter="administrators"
+        )
+    ]
 
 
 async def current_chat_permissions(chat_id):

wbb/modules/log.py

@@ -3,8 +3,8 @@ from datetime import datetime
 import aiofiles
 from pyrogram import filters
 from pyrogram.types import CallbackQuery, Message
-from pyromod.helpers import ikb
 
+from pyromod.helpers import ikb
 from wbb import SUDOERS, app, log_file
 from wbb.utils.pastebin import paste
 

wbb/modules/nsfw.py

@@ -4,7 +4,7 @@ from pyrogram import filters
 
 from wbb import app, arq
 from wbb.core.decorators.errors import capture_err
-from wbb.modules.admin import member_permissions
+from wbb.core.decorators.permissions import adminsOnly
 from wbb.utils.dbfunctions import is_nsfw_on, nsfw_off, nsfw_on
 from wbb.utils.filter_groups import nsfw_detect_group
 
@@ -146,7 +146,7 @@ async def nsfw_scan_command(_, message):
 
 
 @app.on_message(filters.command("anti_nsfw") & ~filters.private)
-@capture_err
+@adminsOnly("can_change_info")
 async def nsfw_enable_disable(_, message):
     if len(message.command) != 2:
         await message.reply_text("Usage: /anti_nsfw [enable | disable]")
@@ -154,11 +154,6 @@ async def nsfw_enable_disable(_, message):
     status = message.text.split(None, 1)[1].strip()
     status = status.lower()
     chat_id = message.chat.id
-    user_id = message.from_user.id
-    permissions = await member_permissions(chat_id, user_id)
-    if "can_change_info" not in permissions:
-        await message.reply_text("You don't have enough permissions.")
-        return
     if status == "enable":
         await nsfw_on(chat_id)
         await message.reply_text(

wbb/modules/spam.py

@@ -28,6 +28,7 @@ from pyrogram.types import (CallbackQuery, InlineKeyboardButton,
 from wbb import LOG_GROUP_ID, SUDOERS, app, arq
 from wbb.core.decorators.errors import capture_err
 from wbb.modules.admin import list_admins, member_permissions
+from wbb.modules.trust import get_spam_data
 from wbb.utils.filter_groups import spam_protection_group
 
 __MODULE__ = "AntiSpam"
@@ -53,18 +54,16 @@ As of now, you cannot turn this off, but we'll add an enable/disable command in
     group=spam_protection_group,
 )
 async def spam_protection_func(_, message: Message):
-    text = message.text if message.text else message.caption
-    user = message.from_user
+    text = message.text or message.caption
     chat_id = message.chat.id
+    user = message.from_user
     if not text or not user:
         return
 
     # We'll handle admins only if it's spam, ignore only sudo users for now.
     if user.id in SUDOERS:
         return
-
-    data = await arq.nlp(text)
-    data = data.result[0]
+    data = await get_spam_data(message, text)
     if not data.is_spam:
         return
     if user.id in (await list_admins(chat_id)):
@@ -72,7 +71,7 @@ async def spam_protection_func(_, message: Message):
     text = f"""
 🚨  **SPAM DETECTED** 🚨
 
-**User:** {message.from_user.mention}
+**User:** {user.mention}
 **Message:** [Link]({message.link})
 **Spam Probability:** {data.spam_probability} %
 **Action:** Alerted
@@ -82,13 +81,11 @@ async def spam_protection_func(_, message: Message):
             [
                 InlineKeyboardButton(
                     text="Yes it's spam", callback_data="s_p_spam"
-                )
-            ],
-            [
+                ),
                 InlineKeyboardButton(
                     text="No, it's not spam",
                     callback_data="s_p_ham",
-                )
+                ),
             ],
         ]
     )

wbb/modules/trust.py

@@ -21,6 +21,8 @@ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 """
+from asyncio import sleep
+
 from pyrogram import filters
 from pyrogram.types import Message
 
@@ -29,10 +31,22 @@ from wbb import (BOT_ID, MESSAGE_DUMP_CHAT, USERBOT_BOT_CHAT_DIFFERENCE, app,
 from wbb.utils.dbfunctions import get_trust_db, update_trust_db
 from wbb.utils.filter_groups import trust_group
 
+spam_db = {}
+
+
+async def get_spam_data(message: Message, text: str):
+    c, m = message.chat.id, message.message_id
+    if c not in spam_db:
+        spam_db[c] = {}
+    if m not in spam_db[c]:
+        data = (await arq.nlp(text)).result[0]
+        spam_db[c][m] = data
+    return spam_db[c][m]
+
 
 @app2.on_message(
     (filters.text | filters.caption)
-    & ~filters.chat([MESSAGE_DUMP_CHAT, BOT_ID])
+    & ~filters.chat(BOT_ID)
     & filters.chat(USERBOT_BOT_CHAT_DIFFERENCE)
     & ~filters.me,
     group=trust_group,
@@ -45,18 +59,20 @@ from wbb.utils.filter_groups import trust_group
     group=trust_group,
 )
 async def trust_watcher_func(_, message: Message):
+    # Sleeping so that we can get cached data created by
+    # spam.py with get_spam_data function
     if message.command:
         return
     if not message.from_user:
         return
     user_id = message.from_user.id
-    text = message.text if message.text else message.caption
+    text = message.text or message.caption
     text = text.strip()
     if not text:
         return
     if len(text) < 2:
         return
-    data = (await arq.nlp(text)).result[0]
+    data = await get_spam_data(message, text)
     spam = data.spam
     await update_trust_db(user_id, spam)