feat: implement multi-level sudo system and core stability fixes

.env.example

@@ -2,7 +2,8 @@ API_ID=
 API_HASH=
 SESSION=
 BOT_TOKEN=
-DB_URL=
+MONGO_URL=
+REDIS_URL=
 OtherSessions=
 PREFIX=
 GEMINI_API_KEY=

Database/Methods/__init__.py

@@ -1,6 +1,10 @@
 from .repeatMethods import RepeatMethods
+from .sessionMethods import SessionMethods
+from .afkMethods import AFKMethods
 
 class Methods(
-    RepeatMethods
+    RepeatMethods,
+    SessionMethods,
+    AFKMethods
 ):
     pass

Database/Methods/afkMethods.py

@@ -0,0 +1,20 @@
+from typing import Optional, Dict, Any
+import time
+
+class AFKMethods:
+    async def set_afk(self, user_id: int, reason: str):
+        """Enable AFK status for a user."""
+        await self.db["afk"].update_one(
+            {"user_id": user_id},
+            {"$set": {"reason": reason, "time": time.time()}},
+            upsert=True
+        )
+
+    async def get_afk(self, user_id: int) -> Optional[Dict[str, Any]]:
+        """Retrieve AFK information for a user."""
+        return await self.db["afk"].find_one({"user_id": user_id})
+
+    async def remove_afk(self, user_id: int) -> bool:
+        """Disable AFK status for a user."""
+        result = await self.db["afk"].delete_one({"user_id": user_id})
+        return result.deleted_count > 0

Database/Methods/repeatMethods.py

@@ -1,102 +1,92 @@
-from sqlalchemy import select, delete
-from Database.Tables.repeatMessage import RepeatMessage
-from Database.Tables.repeatMessageGroup import RepeatMessageGroup
-from Database.Tables.repeatMessageGroupChat import RepeatMessageGroupChat
-
-
 class RepeatMethods:
     # ---------- Groups ----------
 
-    async def create_group(self, name: str, user_id: int) -> RepeatMessageGroup:
-        name = (name.replace(' ', '')).lower() # remove spaces and change to lower case
-        async with self.get_pg() as session: # type: ignore
-            group = RepeatMessageGroup(userId=user_id, name=name)
-            session.add(group)
-            await session.commit()
-            await session.refresh(group)
-            return group
+    async def create_group(self, name: str, user_id: int):
+        name = (name.replace(' ', '')).lower()
+        group = {
+            "userId": user_id,
+            "name": name
+        }
+        result = await self.repeat_groups.insert_one(group)
+        group["_id"] = result.inserted_id
+        return group
 
-    async def get_group(self, group_id: int, user_id: int):
-        async with self.get_pg() as session: # type: ignore
-            x = await session.get(RepeatMessageGroup, group_id)
-            if x.userId == user_id: # Check if the group belongs to the user.
-                return x
+    async def get_group(self, group_id: str, user_id: int):
+        # Note: Using string IDs for Mongo compatibility or ObjectId if needed
+        from bson import ObjectId
+        try:
+            gid = ObjectId(group_id) if isinstance(group_id, str) else group_id
+        except:
+            return None
+        x = await self.repeat_groups.find_one({"_id": gid})
+        if x and x.get("userId") == user_id:
+            return x
 
     async def get_group_by_name(self, name: str, user_id: int):
-        name = (name.replace(' ', '')).lower() # remove spaces and change to lower case
-        async with self.get_pg() as session: # type: ignore
-            q = await session.execute(
-                select(RepeatMessageGroup).where(
-                    RepeatMessageGroup.name == name
-                )
-            )
-            x = q.scalar_one_or_none()
-            if x and x.userId == user_id: # Check if the group belongs to the user.
-                return x
+        name = (name.replace(' ', '')).lower()
+        x = await self.repeat_groups.find_one({"name": name, "userId": user_id})
+        return x
 
     async def get_groups(self, user_id: int):
-        async with self.get_pg() as session:  # type: ignore
-            q = await session.execute(
-                select(RepeatMessageGroup)
-                .where(RepeatMessageGroup.userId == user_id)
-            )
-            return q.scalars().all()
+        cursor = self.repeat_groups.find({"userId": user_id})
+        return await cursor.to_list(length=None)
 
-    async def delete_group(self, group_id: int, user_id: int):
+    async def delete_group(self, group_id: str, user_id: int):
+        from bson import ObjectId
+        try:
+            gid = ObjectId(group_id) if isinstance(group_id, str) else group_id
+        except:
+            raise Exception('Invalid Group ID')
+
         group = await self.get_group(group_id, user_id)
         if not group:
             raise Exception('The group is not found or it does not belongs to you')
-        async with self.get_pg() as session: # type: ignore
-            await session.execute(
-                delete(RepeatMessageGroupChat)
-                .where(RepeatMessageGroupChat.group_id == group_id)
-            )
-            await session.execute(
-                delete(RepeatMessageGroup)
-                .where(RepeatMessageGroup.id == group_id)
-            )
-            await session.commit()
+
+        await self.repeat_group_chats.delete_many({"group_id": gid})
+        await self.repeat_groups.delete_one({"_id": gid})
 
     # ---------- Group Chats ----------
 
-    async def add_chat_to_group(self, group_id: int, chat_id: int, user_id: int):
+    async def add_chat_to_group(self, group_id: str, chat_id: int, user_id: int):
+        from bson import ObjectId
+        gid = ObjectId(group_id) if isinstance(group_id, str) else group_id
+
         group = await self.get_group(group_id, user_id)
         if not group:
             raise Exception('The group is not found or it does not belongs to you')
-        async with self.get_pg() as session: # type: ignore
-            row = RepeatMessageGroupChat(
-                group_id=group_id,
-                chat_id=chat_id,
-                userId=user_id
-            )
-            session.add(row)
-            await session.commit()
-            return row
+
+        row = {
+            "group_id": gid,
+            "chat_id": chat_id,
+            "userId": user_id
+        }
+        await self.repeat_group_chats.insert_one(row)
+        return row
+
+    async def remove_chat_from_group(self, group_id: str, chat_id: int, user_id: int):
+        from bson import ObjectId
+        gid = ObjectId(group_id) if isinstance(group_id, str) else group_id
 
-    async def remove_chat_from_group(self, group_id: int, chat_id: int, user_id: int):
         group = await self.get_group(group_id, user_id)
         if not group:
             raise Exception('The group is not found or it does not belongs to you')
-        async with self.get_pg() as session: # type: ignore
-            await session.execute(
-                delete(RepeatMessageGroupChat)
-                .where(
-                    RepeatMessageGroupChat.group_id == group_id,
-                    RepeatMessageGroupChat.chat_id == chat_id
-                )
-            )
-            await session.commit()
+            
+        await self.repeat_group_chats.delete_one({
+            "group_id": gid,
+            "chat_id": chat_id
+        })
+
+    async def get_group_chats(self, group_id: str, user_id: int) -> list[int]:
+        from bson import ObjectId
+        gid = ObjectId(group_id) if isinstance(group_id, str) else group_id
 
-    async def get_group_chats(self, group_id: int, user_id: int) -> list[int]:
         group = await self.get_group(group_id, user_id)
         if not group:
             raise Exception('The group is not found or it does not belongs to you')
-        async with self.get_pg() as session: # type: ignore
-            q = await session.execute(
-                select(RepeatMessageGroupChat.chat_id)
-                .where(RepeatMessageGroupChat.group_id == group_id)
-            )
-            return [x[0] for x in q.all()]
+
+        cursor = self.repeat_group_chats.find({"group_id": gid})
+        chats = await cursor.to_list(length=None)
+        return [x["chat_id"] for x in chats]
 
     # ---------- Repeat Messages ----------
 
@@ -106,32 +96,27 @@ async def create_repeat_message(
         userId: int,
         message_id: int,
         source_chat_id: int,
-        group_id: int
+        group_id: str
     ):
-        async with self.get_pg() as session: # type: ignore
-            row = RepeatMessage(
-                repeatTime=repeatTime,
-                userId=userId,
-                message_id=message_id,
-                source_chat_id=source_chat_id,
-                group_id=group_id
-            )
-            session.add(row)
-            await session.commit()
-            await session.refresh(row)
-            return row
+        from bson import ObjectId
+        gid = ObjectId(group_id) if isinstance(group_id, str) else group_id
+
+        row = {
+            "repeatTime": repeatTime,
+            "userId": userId,
+            "message_id": message_id,
+            "source_chat_id": source_chat_id,
+            "group_id": gid
+        }
+        result = await self.repeat_messages.insert_one(row)
+        row["_id"] = result.inserted_id
+        return row
 
-    async def get_repeat_messages(self) -> list[RepeatMessage]:
-        async with self.get_pg() as session: # type: ignore
-            q = await session.execute(
-                select(RepeatMessage)
-            )
-            return q.scalars().all()
+    async def get_repeat_messages(self):
+        cursor = self.repeat_messages.find({})
+        return await cursor.to_list(length=None)
 
-    async def delete_repeat_message(self, repeat_id: int):
-        async with self.get_pg() as session: # type: ignore
-            await session.execute(
-                delete(RepeatMessage)
-                .where(RepeatMessage.id == repeat_id)
-            )
-            await session.commit()
+    async def delete_repeat_message(self, repeat_id: str):
+        from bson import ObjectId
+        rid = ObjectId(repeat_id) if isinstance(repeat_id, str) else repeat_id
+        await self.repeat_messages.delete_one({"_id": rid})

Database/Methods/sessionMethods.py

@@ -0,0 +1,21 @@
+from typing import List, Optional
+
+class SessionMethods:
+    async def add_session(self, session_string: str) -> bool:
+        """Add a new session string to the database if it doesn't already exist."""
+        existing = await self.db["sessions"].find_one({"session": session_string})
+        if not existing:
+            await self.db["sessions"].insert_one({"session": session_string})
+            return True
+        return False
+
+    async def get_all_sessions(self) -> List[str]:
+        """Retrieve all extra session strings from the database."""
+        cursor = self.db["sessions"].find({})
+        sessions = await cursor.to_list(length=100)
+        return [s["session"] for s in sessions]
+
+    async def remove_session(self, session_string: str) -> bool:
+        """Remove a session string from the database."""
+        result = await self.db["sessions"].delete_one({"session": session_string})
+        return result.deleted_count > 0

Database/Methods/sudoMethods.py

@@ -0,0 +1,88 @@
+from typing import List, Optional
+import logging
+
+logger = logging.getLogger("Hazel.SudoMethods")
+
+class SudoMethods:
+    # MongoDB collections are initialized in MongoClient
+    # We use Redis for fast lookup: "sudo_list" and "fsudo_list" (Sets)
+
+    async def add_sudo(self, user_id: int, level: str = "sudo"):
+        """
+        level can be 'sudo' or 'fsudo'
+        """
+        user_id = int(user_id)
+        # 1. Store in MongoDB
+        collection = self.db["sudo_users"]
+        await collection.update_one(
+            {"user_id": user_id},
+            {"$set": {"level": level}},
+            upsert=True
+        )
+
+        # 2. Update Redis
+        from Hazel import Redis
+        if Redis:
+            set_name = f"{level}_list"
+            await Redis.redis.sadd(set_name, str(user_id))
+
+        logger.info(f"Added user {user_id} as {level}")
+
+    async def remove_sudo(self, user_id: int):
+        user_id = int(user_id)
+        # 1. Remove from MongoDB
+        collection = self.db["sudo_users"]
+        await collection.delete_one({"user_id": user_id})
+
+        # 2. Remove from Redis
+        from Hazel import Redis
+        if Redis:
+            await Redis.redis.srem("sudo_list", str(user_id))
+            await Redis.redis.srem("fsudo_list", str(user_id))
+
+        logger.info(f"Removed user {user_id} from sudo list")
+
+    async def get_all_sudo(self) -> List[dict]:
+        collection = self.db["sudo_users"]
+        return await collection.find({}).to_list(length=None)
+
+    async def is_sudo(self, user_id: int) -> bool:
+        user_id = int(user_id)
+        from Hazel import Redis
+        if Redis:
+            # Check both as strings (we store them as strings in Redis)
+            res1 = await Redis.redis.sismember("sudo_list", str(user_id))
+            res2 = await Redis.redis.sismember("fsudo_list", str(user_id))
+            is_it = bool(res1 or res2)
+            return is_it
+
+        collection = self.db["sudo_users"]
+        res = await collection.find_one({"user_id": user_id})
+        return res is not None
+
+    async def is_fsudo(self, user_id: int) -> bool:
+        user_id = int(user_id)
+        from Hazel import Redis
+        if Redis:
+            res = await Redis.redis.sismember("fsudo_list", str(user_id))
+            return bool(res)
+
+        collection = self.db["sudo_users"]
+        res = await collection.find_one({"user_id": user_id, "level": "fsudo"})
+        return res is not None
+
+    async def reload_sudo_cache(self):
+        """Load all sudo users from Mongo into Redis."""
+        from Hazel import Redis
+        if not Redis: return
+
+        # Clear existing
+        await Redis.redis.delete("sudo_list")
+        await Redis.redis.delete("fsudo_list")
+
+        all_users = await self.get_all_sudo()
+        for u in all_users:
+            level = u.get("level", "sudo")
+            await Redis.redis.sadd(f"{level}_list", str(u["user_id"]))
+
+        logger.info(f"Reloaded {len(all_users)} sudo users into Redis cache")