Implement user ban system

bot.py

@@ -3,7 +3,8 @@
 from datetime import datetime
 from pyrogram import Client
 from pyrogram.enums import ParseMode
-from config import API_HASH, APP_ID, LOGGER, TG_BOT_TOKEN, TG_BOT_WORKERS, PORT
+from config import API_HASH, APP_ID, LOGGER, TG_BOT_TOKEN, TG_BOT_WORKERS, PORT, ADMINS
+from database.database import add_admin
 from plugins import web_server
 import pyrogram.utils
 from aiohttp import web
@@ -35,6 +36,13 @@ async def start(self):
         self.LOGGER(__name__).info(f"{name}")
         self.username = usr_bot_me.username
 
+        # Seed DB admins from config at startup
+        try:
+            for admin_id in ADMINS:
+                await add_admin(int(admin_id))
+        except Exception as e:
+            self.LOGGER(__name__).warning(f"Failed to sync ADMINS to DB: {e}")
+
         # Web-response
         app = web.AppRunner(await web_server())
         await app.setup()

database/database.py

@@ -2,13 +2,16 @@
 import base64
 from config import DB_URI, DB_NAME
 from datetime import datetime, timedelta
+from typing import List
 
 dbclient = motor.motor_asyncio.AsyncIOMotorClient(DB_URI)
 database = dbclient[DB_NAME]
 
 user_data = database['users']
 channels_collection = database["channels"]
 encoded_links_collection = database["links"]
+admins_collection = database["admins"]
+banned_collection = database["banned_users"]
 
 async def add_user(user_id: int):
     existing_user = await user_data.find_one({'_id': user_id})
@@ -33,7 +36,20 @@ async def del_user(user_id: int):
 
 ##-------------------------------------------------------------------
 
-async def is_admin(user_id: int):
+async def add_admin(user_id: int) -> None:
+    await admins_collection.update_one({'_id': user_id}, {'$set': {'_id': user_id}}, upsert=True)
+
+
+async def remove_admin(user_id: int) -> None:
+    await admins_collection.delete_one({'_id': user_id})
+
+
+async def list_admins() -> List[int]:
+    cursor = admins_collection.find({}, {'_id': 1})
+    return [doc['_id'] async for doc in cursor]
+
+
+async def is_admin(user_id: int) -> bool:
     return bool(await admins_collection.find_one({'_id': user_id}))
 
 ##-------------------------------------------------------------------
@@ -87,3 +103,22 @@ async def get_channel_by_encoded_link2(encoded_link: str):
     return channel["channel_id"] if channel else None
 
 
+##-------------------------------------------------------------------
+## Ban system
+
+async def add_ban(user_id: int) -> None:
+    await banned_collection.update_one({'_id': user_id}, {'$set': {'_id': user_id, 'banned_at': datetime.utcnow()}}, upsert=True)
+
+
+async def remove_ban(user_id: int) -> None:
+    await banned_collection.delete_one({'_id': user_id})
+
+
+async def is_banned(user_id: int) -> bool:
+    return bool(await banned_collection.find_one({'_id': user_id}))
+
+
+async def list_banned() -> List[int]:
+    cursor = banned_collection.find({}, {'_id': 1})
+    return [doc['_id'] async for doc in cursor]
+

plugins/admin_mgmt.py

@@ -0,0 +1,49 @@
+from pyrogram import Client as Bot, filters
+from pyrogram.types import Message
+from config import OWNER_ID
+from database.database import add_admin, remove_admin, list_admins, is_admin, add_ban
+
+
+@Bot.on_message(filters.command('addadmin') & filters.private & filters.user(OWNER_ID))
+async def add_admin_cmd(_: Bot, message: Message):
+    try:
+        user_id = int(message.command[1])
+    except (IndexError, ValueError):
+        return await message.reply("Usage: /addadmin <user_id>")
+
+    await add_admin(user_id)
+    return await message.reply(f"✅ Added admin: {user_id}")
+
+
+@Bot.on_message(filters.command('deladmin') & filters.private & filters.user(OWNER_ID))
+async def del_admin_cmd(_: Bot, message: Message):
+    try:
+        user_id = int(message.command[1])
+    except (IndexError, ValueError):
+        return await message.reply("Usage: /deladmin <user_id>")
+
+    await remove_admin(user_id)
+    return await message.reply(f"❌ Removed admin: {user_id}")
+
+
+@Bot.on_message(filters.command('admins') & filters.private & filters.user(OWNER_ID))
+async def list_admins_cmd(_: Bot, message: Message):
+    admins = await list_admins()
+    if not admins:
+        return await message.reply("No admins set.")
+    return await message.reply("Admins:\n" + "\n".join([str(a) for a in admins]))
+
+
+@Bot.on_message(filters.command('ban') & filters.private)
+async def ban_cmd(_: Bot, message: Message):
+    requester_id = message.from_user.id
+    if not (requester_id == OWNER_ID or await is_admin(requester_id)):
+        return
+
+    try:
+        user_id = int(message.command[1])
+    except (IndexError, ValueError):
+        return await message.reply("Usage: /ban <user_id>")
+
+    await add_ban(user_id)
+    return await message.reply(f"🚫 Banned: {user_id}")

plugins/newpost.py

@@ -4,8 +4,8 @@
 from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton, Message
 from pyrogram.errors import UserNotParticipant, FloodWait, ChatAdminRequired, RPCError
 from pyrogram.errors import InviteHashExpired, InviteRequestSent
-from database.database import save_channel, delete_channel, get_channels
-from config import ADMINS, OWNER_ID
+from database.database import save_channel, delete_channel, get_channels, is_admin
+from config import OWNER_ID
 from database.database import save_encoded_link, get_channel_by_encoded_link, save_encoded_link2, get_channel_by_encoded_link2
 from helper_func import encode
 from datetime import datetime, timedelta
@@ -27,10 +27,14 @@ async def revoke_invite_after_10_minutes(client: Bot, channel_id: int, link: str
 ##----------------------------------------------------------------------------------------------------        
 ##----------------------------------------------------------------------------------------------------
 
-@Bot.on_message(filters.command('setchannel') & filters.private & filters.user(OWNER_ID))
+@Bot.on_message(filters.command('setchannel') & filters.private)
 async def set_channel(client: Bot, message: Message):
     user_id = message.from_user.id
-    if user_id not in ADMINS:
+    # Ignore if banned
+    from database.database import is_banned
+    if await is_banned(user_id):
+        return
+    if not (user_id == OWNER_ID or await is_admin(user_id)):
         return await message.reply("You are not an admin.")
 
     try:
@@ -60,10 +64,13 @@ async def set_channel(client: Bot, message: Message):
 ##----------------------------------------------------------------------------------------------------
 ##----------------------------------------------------------------------------------------------------        
 
-@Bot.on_message(filters.command('delchannel') & filters.private & filters.user(OWNER_ID))
+@Bot.on_message(filters.command('delchannel') & filters.private)
 async def del_channel(client: Bot, message: Message):
     user_id = message.from_user.id
-    if user_id not in ADMINS:
+    from database.database import is_banned
+    if await is_banned(user_id):
+        return
+    if not (user_id == OWNER_ID or await is_admin(user_id)):
         return await message.reply("You are not an admin.")
 
     try:
@@ -77,8 +84,14 @@ async def del_channel(client: Bot, message: Message):
 ##----------------------------------------------------------------------------------------------------        
 ##----------------------------------------------------------------------------------------------------
 
-@Bot.on_message(filters.command('channelpost') & filters.private & filters.user(OWNER_ID))
+@Bot.on_message(filters.command('channelpost') & filters.private)
 async def channel_post(client: Bot, message: Message):
+    user_id = message.from_user.id
+    from database.database import is_banned
+    if await is_banned(user_id):
+        return
+    if not (user_id == OWNER_ID or await is_admin(user_id)):
+        return await message.reply("You are not an admin.")
     channels = await get_channels()
     if not channels:
         return await message.reply("No channels available. Use /setchannel first.")
@@ -104,8 +117,14 @@ async def channel_post(client: Bot, message: Message):
 #-------------------------------------------------------------------------------------------------------------------------
 #------------------------------------------------------------------------------------------------------------------------- 
 
-@Bot.on_message(filters.command('reqpost') & filters.private & filters.user(OWNER_ID))
+@Bot.on_message(filters.command('reqpost') & filters.private)
 async def req_post(client: Bot, message: Message):
+    user_id = message.from_user.id
+    from database.database import is_banned
+    if await is_banned(user_id):
+        return
+    if not (user_id == OWNER_ID or await is_admin(user_id)):
+        return await message.reply("You are not an admin.")
     channels = await get_channels()
     if not channels:
         return await message.reply("No channels available. Use /setchannel first.")

plugins/start.py

@@ -11,10 +11,10 @@
 
 from bot import Bot
 from datetime import datetime, timedelta
-from config import ADMINS, OWNER_ID
+from config import OWNER_ID
 from helper_func import encode, decode
 from database.database import save_encoded_link, get_channel_by_encoded_link, save_encoded_link2, get_channel_by_encoded_link2
-from database.database import add_user, del_user, full_userbase, present_user, is_admin
+from database.database import add_user, del_user, full_userbase, present_user, is_admin, is_banned
 from plugins.newpost import revoke_invite_after_10_minutes
 
 #=====================================================================================##
@@ -23,12 +23,13 @@
 async def start_command(client: Bot, message: Message):
     user_id = message.from_user.id
 
-    # Check if the user is banned
-    if user_id in user_banned_until:
-        # Check if the ban duration has not expired
-        if datetime.now() < user_banned_until[user_id]:
-            # User is still banned, do not process the command
-            return await message.reply_text("🚫 You are temporarily banned from using commands due to spamming. Try again later.")
+    # Ignore if user is permanently banned
+    if await is_banned(user_id):
+        return
+
+    # Check if the user is temporarily banned
+    if user_id in user_banned_until and datetime.now() < user_banned_until[user_id]:
+        return
 
     # Proceed with the original functionality if the user is not banned
     text = message.text
@@ -147,6 +148,9 @@ async def send_text(client: Bot, message: Message):
 
 @Bot.on_callback_query(filters.regex("help"))
 async def help_callback(client: Bot, callback_query):
+    # Ignore permanently banned users
+    if await is_banned(callback_query.from_user.id):
+        return
     # Define the inline keyboard with the "Close" button
     inline_buttons = InlineKeyboardMarkup(
         [
@@ -160,6 +164,8 @@ async def help_callback(client: Bot, callback_query):
 
 @Bot.on_callback_query(filters.regex("close"))
 async def close_callback(client: Bot, callback_query):
+    if await is_banned(callback_query.from_user.id):
+        return
     await callback_query.answer()
     await callback_query.message.delete()
 
@@ -178,12 +184,12 @@ async def monitor_messages(client: Bot, message: Message):
     user_id = message.from_user.id
     now = datetime.now()
 
-    if user_id in ADMINS:
-        return 
+    # Ignore for owner/admins
+    if user_id == OWNER_ID or await is_admin(user_id):
+        return
 
-
-    if user_id in user_banned_until and now < user_banned_until[user_id]:
-        await message.reply_text("⚠️ You are temporarily banned from using commands due to spamming. Try again later.")
+    # Ignore permanently banned users entirely
+    if await is_banned(user_id):
         return
 
     if user_id not in user_message_count:
@@ -195,5 +201,4 @@ async def monitor_messages(client: Bot, message: Message):
 
     if len(user_message_count[user_id]) > MAX_MESSAGES:
         user_banned_until[user_id] = now + BAN_DURATION
-        await message.reply_text("🚫 You have been temporarily banned for spamming. Try again in 1 hour.")
         return