Create database.py

database.py

@@ -0,0 +1,676 @@
+import datetime
+import time
+
+from motor import motor_asyncio
+from motor.core import AgnosticClient
+
+from Akeno.utils.logger import LOGS
+from config import MONGO_URL
+
+
+class Database:
+    def __init__(self, uri: str) -> None:
+        self.client: AgnosticClient = motor_asyncio.AsyncIOMotorClient(uri)
+        self.db = self.client["Akeno"]
+
+        self.afk = self.db["afk"]
+        self.antiflood = self.db["antiflood"]
+        self.autopost = self.db["autopost"]
+        self.blacklist = self.db["blacklist"]
+        self.echo = self.db["echo"]
+        self.env = self.db["env"]
+        self.filter = self.db["filter"]
+        self.forcesub = self.db["forcesub"]
+        self.gachabots = self.db["gachabots"]
+        self.cohere = self.db["cohere"]
+        self.chatbot = self.db["chatbot"]
+        self.backup_chatbot = self.db["backupchatbot"]
+        self.antiarabic = self.db["antiarabic"]
+        self.antinsfw = self.db["antinsfw"]
+        self.gban = self.db["gban"]
+        self.gmute = self.db["gmute"]
+        self.greetings = self.db["greetings"]
+        self.mute = self.db["mute"]
+        self.pmpermit = self.db["pmpermit"]
+        self.session = self.db["session"]
+        self.snips = self.db["snips"]
+        self.stan_users = self.db["stan_users"]
+
+    async def connect(self):
+        try:
+            await self.client.admin.command("ping")
+            LOGS.info(f"Database Connection Established!")
+        except Exception as e:
+            LOGS.info(f"DatabaseErr: {e} ")
+            quit(1)
+
+    async def _close(self):
+        await self.client.close()
+
+    def get_datetime(self) -> str:
+        return datetime.datetime.now().strftime("%d/%m/%Y - %H:%M")
+
+    async def set_env(self, name: str, value: str) -> None:
+        await self.env.update_one(
+            {"name": name}, {"$set": {"value": value}}, upsert=True
+        )
+
+    async def get_env(self, name: str) -> str | None:
+        if await self.is_env(name):
+            data = await self.env.find_one({"name": name})
+            return data["value"]
+        return None
+
+    async def rm_env(self, name: str) -> None:
+        await self.env.delete_one({"name": name})
+
+    async def is_env(self, name: str) -> bool:
+        if await self.env.find_one({"name": name}):
+            return True
+        return False
+
+    async def get_all_env(self) -> list:
+        return [i async for i in self.env.find({})]
+
+    async def is_stan(self, client: int, user_id: int) -> bool:
+        if await self.stan_users.find_one({"client": client, "user_id": user_id}):
+            return True
+        return False
+
+    async def add_stan(self, client: int, user_id: int) -> bool:
+        if await self.is_stan(client, user_id):
+            return False
+        await self.stan_users.insert_one(
+            {"client": client, "user_id": user_id, "date": self.get_datetime()}
+        )
+        return True
+
+    async def rm_stan(self, client: int, user_id: int) -> bool:
+        if not await self.is_stan(client, user_id):
+            return False
+        await self.stan_users.delete_one({"client": client, "user_id": user_id})
+        return True
+
+    async def get_stans(self, client: int) -> list:
+        return [i async for i in self.stan_users.find({"client": client})]
+
+    async def get_all_stans(self) -> list:
+        return [i async for i in self.stan_users.find({})]
+
+    async def is_session(self, user_id: int) -> bool:
+        if await self.session.find_one({"user_id": user_id}):
+            return True
+        return False
+
+    async def update_session(self, user_id: int, session: str) -> None:
+        await self.session.update_one(
+            {"user_id": user_id},
+            {"$set": {"session": session, "date": self.get_datetime()}},
+            upsert=True,
+        )
+
+    async def rm_session(self, user_id: int) -> None:
+        await self.session.delete_one({"user_id": user_id})
+
+    async def get_session(self, user_id: int):
+        if not await self.is_session(user_id):
+            return False
+        data = await self.session.find_one({"user_id": user_id})
+        return data
+
+    async def get_all_sessions(self) -> list:
+        return [i async for i in self.session.find({})]
+
+    async def is_gbanned(self, user_id: int) -> bool:
+        if await self.gban.find_one({"user_id": user_id}):
+            return True
+        return False
+
+    async def add_gban(self, user_id: int, reason: str) -> bool:
+        if await self.is_gbanned(user_id):
+            return False
+        await self.gban.insert_one(
+            {"user_id": user_id, "reason": reason, "date": self.get_datetime()}
+        )
+        return True
+
+    async def rm_gban(self, user_id: int):
+        if not await self.is_gbanned(user_id):
+            return None
+        reason = (await self.gban.find_one({"user_id": user_id}))["reason"]
+        await self.gban.delete_one({"user_id": user_id})
+        return reason
+
+    async def get_gban(self) -> list:
+        return [i async for i in self.gban.find({})]
+
+    async def get_gban_user(self, user_id: int) -> dict | None:
+        if not await self.is_gbanned(user_id):
+            return None
+        return await self.gban.find_one({"user_id": user_id})
+
+    async def is_gmuted(self, user_id: int) -> bool:
+        if await self.gmute.find_one({"user_id": user_id}):
+            return True
+        return False
+
+    async def add_gmute(self, user_id: int, reason: str) -> bool:
+        if await self.is_gmuted(user_id):
+            return False
+        await self.gmute.insert_one(
+            {"user_id": user_id, "reason": reason, "date": self.get_datetime()}
+        )
+        return True
+
+    async def rm_gmute(self, user_id: int):
+        if not await self.is_gmuted(user_id):
+            return None
+        reason = (await self.gmute.find_one({"user_id": user_id}))["reason"]
+        await self.gmute.delete_one({"user_id": user_id})
+        return reason
+
+    async def get_gmute(self) -> list:
+        return [i async for i in self.gmute.find({})]
+
+    async def add_mute(self, client: int, user_id: int, chat_id: int, reason: str):
+        await self.mute.update_one(
+            {"client": client, "user_id": user_id, "chat_id": chat_id},
+            {"$set": {"reason": reason, "date": self.get_datetime()}},
+            upsert=True,
+        )
+
+    async def rm_mute(self, client: int, user_id: int, chat_id: int) -> str:
+        reason = (await self.get_mute(client, user_id, chat_id))["reason"]
+        await self.mute.delete_one({"client": client, "user_id": user_id, "chat_id": chat_id})
+        return reason
+
+    async def is_muted(self, client: int, user_id: int, chat_id: int) -> bool:
+        if await self.get_mute(client, user_id, chat_id):
+            return True
+        return False
+
+    async def get_mute(self, client: int, user_id: int, chat_id: int):
+        data = await self.mute.find_one({"client": client, "user_id": user_id, "chat_id": chat_id})
+        return data
+
+    async def set_afk(
+        self, user_id: int, reason: str, media: int, media_type: str
+    ) -> None:
+        await self.afk.update_one(
+            {"user_id": user_id},
+            {
+                "$set": {
+                    "reason": reason,
+                    "time": time.time(),
+                    "media": media,
+                    "media_type": media_type,
+                }
+            },
+            upsert=True,
+        )
+
+    async def get_afk(self, user_id: int):
+        data = await self.afk.find_one({"user_id": user_id})
+        return data
+
+    async def is_afk(self, user_id: int) -> bool:
+        if await self.afk.find_one({"user_id": user_id}):
+            return True
+        return False
+
+    async def rm_afk(self, user_id: int) -> None:
+        await self.afk.delete_one({"user_id": user_id})
+
+    async def set_flood(self, client_chat: tuple[int, int], settings: dict):
+        await self.antiflood.update_one(
+            {"client": client_chat[0], "chat": client_chat[1]},
+            {"$set": settings},
+            upsert=True,
+        )
+
+    async def get_flood(self, client_chat: tuple[int, int]):
+        data = await self.antiflood.find_one(
+            {"client": client_chat[0], "chat": client_chat[1]}
+        )
+        return data or {}
+
+    async def is_flood(self, client_chat: tuple[int, int]) -> bool:
+        data = await self.get_flood(client_chat)
+
+        if not data:
+            return False
+
+        if data["limit"] == 0:
+            return False
+
+        return True
+
+    async def get_all_floods(self) -> list:
+        return [i async for i in self.antiflood.find({})]
+
+    async def set_autopost(self, client: int, from_channel: int, to_channel: int):
+        await self.autopost.update_one(
+            {"client": client},
+            {
+                "$push": {
+                    "autopost": {
+                        "from_channel": from_channel,
+                        "to_channel": to_channel,
+                        "date": self.get_datetime(),
+                    }
+                }
+            },
+            upsert=True,
+        )
+
+    async def get_autopost(self, client: int, from_channel: int):
+        data = await self.autopost.find_one(
+            {
+                "client": client,
+                "autopost": {"$elemMatch": {"from_channel": from_channel}},
+            }
+        )
+        return data
+
+    async def is_autopost(
+        self, client: int, from_channel: int, to_channel: int = None
+    ) -> bool:
+        if to_channel:
+            data = await self.autopost.find_one(
+                {
+                    "client": client,
+                    "autopost": {
+                        "$elemMatch": {
+                            "from_channel": from_channel,
+                            "to_channel": to_channel,
+                        }
+                    },
+                }
+            )
+        else:
+            data = await self.autopost.find_one(
+                {
+                    "client": client,
+                    "autopost": {"$elemMatch": {"from_channel": from_channel}},
+                }
+            )
+        return True if data else False
+
+    async def rm_autopost(self, client: int, from_channel: int, to_channel: int):
+        await self.autopost.update_one(
+            {"client": client},
+            {
+                "$pull": {
+                    "autopost": {
+                        "from_channel": from_channel,
+                        "to_channel": to_channel,
+                    }
+                }
+            },
+        )
+
+    async def get_all_autoposts(self, client: int) -> list:
+        return [i async for i in self.autopost.find({"client": client})]
+
+    async def add_blacklist(self, client: int, chat: int, blacklist: str):
+        await self.blacklist.update_one(
+            {"client": client, "chat": chat},
+            {"$push": {"blacklist": blacklist}},
+            upsert=True,
+        )
+
+    async def rm_blacklist(self, client: int, chat: int, blacklist: str):
+        await self.blacklist.update_one(
+            {"client": client, "chat": chat},
+            {"$pull": {"blacklist": blacklist}},
+        )
+
+    async def is_blacklist(self, client: int, chat: int, blacklist: str) -> bool:
+        blacklists = await self.get_all_blacklists(client, chat)
+        if blacklist in blacklists:
+            return True
+        return False
+
+    async def get_all_blacklists(self, client: int, chat: int) -> list:
+        data = await self.blacklist.find_one({"client": client, "chat": chat})
+
+        if not data:
+            return []
+
+        return data["blacklist"]
+
+    async def get_blacklist_clients(self) -> list:
+        return [i async for i in self.blacklist.find({})]
+
+    async def set_echo(self, client: int, chat: int, user: int):
+        await self.echo.update_one(
+            {"client": client, "chat": chat},
+            {"$push": {"echo": user}},
+            upsert=True,
+        )
+
+    async def rm_echo(self, client: int, chat: int, user: int):
+        await self.echo.update_one(
+            {"client": client, "chat": chat},
+            {"$pull": {"echo": user}},
+        )
+
+    async def is_echo(self, client: int, chat: int, user: int) -> bool:
+        data = await self.get_all_echo(client, chat)
+        if user in data:
+            return True
+        return False
+
+    async def get_all_echo(self, client: int, chat: int) -> list:
+        data = await self.echo.find_one({"client": client, "chat": chat})
+
+        if not data:
+            return []
+
+        return data["echo"]
+
+    async def set_filter(self, client: int, chat: int, keyword: str, msgid: int):
+        await self.filter.update_one(
+            {"client": client, "chat": chat},
+            {"$push": {"filter": {"keyword": keyword, "msgid": msgid}}},
+            upsert=True,
+        )
+
+    async def rm_filter(self, client: int, chat: int, keyword: str):
+        await self.filter.update_one(
+            {"client": client, "chat": chat},
+            {"$pull": {"filter": {"keyword": keyword}}},
+        )
+
+    async def rm_all_filters(self, client: int, chat: int):
+        await self.filter.delete_one({"client": client, "chat": chat})
+
+    async def is_filter(self, client: int, chat: int, keyword: str) -> bool:
+        data = await self.get_filter(client, chat, keyword)
+        return True if data else False
+
+    async def get_filter(self, client: int, chat: int, keyword: str):
+        data = await self.filter.find_one(
+            {
+                "client": client,
+                "chat": chat,
+                "filter": {"$elemMatch": {"keyword": keyword}},
+            }
+        )
+        return data
+
+    async def get_all_filters(self, client: int, chat: int) -> list:
+        data = await self.filter.find_one({"client": client, "chat": chat})
+
+        if not data:
+            return []
+
+        return data["filter"]
+
+    async def set_snip(self, client: int, chat: int, keyword: str, msgid: int):
+        await self.snips.update_one(
+            {"client": client, "chat": chat},
+            {"$push": {"snips": {"keyword": keyword, "msgid": msgid}}},
+            upsert=True,
+        )
+
+    async def rm_snip(self, client: int, chat: int, keyword: str):
+        await self.snips.update_one(
+            {"client": client, "chat": chat},
+            {"$pull": {"snips": {"keyword": keyword}}},
+        )
+
+    async def rm_all_snips(self, client: int, chat: int):
+        await self.snips.delete_one({"client": client, "chat": chat})
+
+    async def is_snip(self, client: int, chat: int, keyword: str) -> bool:
+        data = await self.get_snip(client, chat, keyword)
+        return True if data else False
+
+    async def get_snip(self, client: int, chat: int, keyword: str):
+        data = await self.snips.find_one(
+            {
+                "client": client,
+                "chat": chat,
+                "snips": {"$elemMatch": {"keyword": keyword}},
+            }
+        )
+        return data
+
+    async def get_all_snips(self, client: int, chat: int) -> list:
+        data = await self.snips.find_one({"client": client, "chat": chat})
+
+        if not data:
+            return []
+
+        return data["snips"]
+
+    async def add_pmpermit(self, client: int, user: int):
+        await self.pmpermit.update_one(
+            {"client": client, "user": user},
+            {"$set": {"date": self.get_datetime()}},
+            upsert=True,
+        )
+
+    async def rm_pmpermit(self, client: int, user: int):
+        await self.pmpermit.delete_one({"client": client, "user": user})
+
+    async def is_pmpermit(self, client: int, user: int) -> bool:
+        data = await self.get_pmpermit(client, user)
+        return True if data else False
+
+    async def get_pmpermit(self, client: int, user: int):
+        data = await self.pmpermit.find_one({"client": client, "user": user})
+        return data
+
+    async def get_all_pmpermits(self, client: int) -> list:
+        return [i async for i in self.pmpermit.find({"client": client})]
+
+    async def set_welcome(self, client: int, chat: int, message: int):
+        await self.greetings.update_one(
+            {"client": client, "chat": chat, "welcome": True},
+            {"$set": {"message": message}},
+            upsert=True,
+        )
+
+    async def rm_welcome(self, client: int, chat: int):
+        await self.greetings.delete_one(
+            {"client": client, "chat": chat, "welcome": True}
+        )
+
+    async def is_welcome(self, client: int, chat: int) -> bool:
+        data = await self.get_welcome(client, chat)
+        return True if data else False
+
+    async def get_welcome(self, client: int, chat: int):
+        data = await self.greetings.find_one(
+            {"client": client, "chat": chat, "welcome": True}
+        )
+        return data
+
+    async def set_goodbye(self, client: int, chat: int, message: int):
+        await self.greetings.update_one(
+            {"client": client, "chat": chat, "welcome": False},
+            {"$set": {"message": message}},
+            upsert=True,
+        )
+
+    async def rm_goodbye(self, client: int, chat: int):
+        await self.greetings.delete_one(
+            {"client": client, "chat": chat, "welcome": False}
+        )
+
+    async def is_goodbye(self, client: int, chat: int) -> bool:
+        data = await self.get_goodbye(client, chat)
+        return True if data else False
+
+    async def get_goodbye(self, client: int, chat: int):
+        data = await self.greetings.find_one(
+            {"client": client, "chat": chat, "welcome": False}
+        )
+        return data
+
+    async def get_all_greetings(self, client: int) -> list:
+        return [i async for i in self.greetings.find({"client": client})]
+
+    async def add_forcesub(self, chat: int, must_join: int):
+        await self.forcesub.update_one(
+            {"chat": chat},
+            {"$push": {"must_join": must_join}},
+            upsert=True,
+        )
+
+    async def rm_forcesub(self, chat: int, must_join: int) -> int:
+        await self.forcesub.update_one(
+            {"chat": chat},
+            {"$pull": {"must_join": must_join}},
+        )
+        data = await self.forcesub.find_one({"chat": chat})
+        return len(data["must_join"])
+
+    async def rm_all_forcesub(self, in_chat: int):
+        await self.forcesub.delete_one({"chat": in_chat})
+
+    async def is_forcesub(self, chat: int, must_join: int) -> bool:
+        data = await self.get_forcesub(chat)
+        if must_join in data["must_join"]:
+            return True
+        return False
+
+    async def get_forcesub(self, in_chat: int):
+        data = await self.forcesub.find_one({"chat": in_chat})
+        return data
+
+    async def get_all_forcesubs(self) -> list:
+        return [i async for i in self.forcesub.find({})]
+
+    async def add_gachabot(
+        self, client: int, bot: tuple[int, str], catch_command: str, chat_id: int
+    ):
+        await self.gachabots.update_one(
+            {"client": client, "bot": bot[0]},
+            {
+                "$set": {
+                    "username": bot[1],
+                    "catch_command": catch_command,
+                    "chat_id": chat_id,
+                    "date": self.get_datetime(),
+                }
+            },
+            upsert=True,
+        )
+
+    async def rm_gachabot(self, client: int, bot: int, chat_id: int = None):
+        if chat_id:
+            await self.gachabots.delete_one(
+                {"client": client, "bot": bot, "chat_id": chat_id}
+            )
+        else:
+            await self.gachabots.delete_one({"client": client, "bot": bot})
+
+    async def is_gachabot(self, client: int, bot: int, chat_id: int) -> bool:
+        data = await self.get_gachabot(client, bot, chat_id)
+        return True if data else False
+
+    async def get_gachabot(self, client: int, bot: int, chat_id: int):
+        data = await self.gachabots.find_one(
+            {"client": client, "bot": bot, "chat_id": chat_id}
+        )
+
+        return data
+
+    async def get_all_gachabots(self, client: int) -> list:
+        return [i async for i in self.gachabots.find({"client": client})]
+
+    async def get_all_gachabots_id(self) -> list:
+        data = await self.gachabots.distinct("bot")
+        return data
+
+    async def _get_cohere_chat_from_db(self, user_id):
+        user_data = await self.cohere.find_one({"user_id": user_id})
+        return user_data.get("cohere_chat", []) if user_data else []
+
+    async def _update_cohere_chat_in_db(self, user_id, cohere_chat):
+        await self.cohere.update_one(
+            {"user_id": user_id},
+            {"$set": {"cohere_chat": cohere_chat}},
+            upsert=True
+        )
+
+    async def _clear_history_in_db(self, user_id):
+        unset_clear = {"cohere_chat": None}
+        return await self.cohere.update_one({"user_id": user_id}, {"$unset": unset_clear})
+
+    async def clear_database(self, user_id):
+        """Clear the cohere history for the current user."""
+        result = await self._clear_history_in_db(user_id)
+        if result.modified_count > 0:
+            return "Chat history cleared successfully."
+        else:
+            return "No chat history found to clear."
+
+    async def set_chat_setting(self, chat_id, isboolean):
+        await self.antiarabic.update_one(
+            {"chat_id": chat_id},
+            {"$set": {"arabic": isboolean}},
+            upsert=True
+        )
+
+    async def chat_antiarabic(self, chat_id):
+        user_data = await self.antiarabic.find_one({"chat_id": chat_id})
+        if user_data:
+            return user_data.get("arabic", False)
+        return False
+
+    async def add_chatbot(self, chat_id, user_id):
+        await self.chatbot.update_one(
+            {"chat_id": chat_id},
+            {"$set": {"user_id": user_id}},
+            upsert=True
+        )
+
+    async def get_chatbot(self, chat_id):
+        user_data = await self.chatbot.find_one({"chat_id": chat_id})
+        return user_data.get("user_id") if user_data else None
+
+    async def remove_chatbot(self, chat_id):
+        unset_data = {"user_id": None}
+        return await self.chatbot.update_one({"chat_id": chat_id}, {"$unset": unset_data})
+
+    async def _update_chatbot_chat_in_db(self, user_id, chatbot_chat):
+        await self.backup_chatbot.update_one(
+            {"user_id": user_id},
+            {"$set": {"chatbot_chat": chatbot_chat}},
+            upsert=True
+        )
+
+    async def _get_chatbot_chat_from_db(self, user_id):
+        user_data = await self.backup_chatbot.find_one({"user_id": user_id})
+        return user_data.get("chatbot_chat", []) if user_data else []
+
+    async def _clear_chatbot_history_in_db(self, user_id):
+        unset_clear = {"chatbot_chat": None}
+        return await self.backup_chatbot.update_one({"user_id": user_id}, {"$unset": unset_clear})
+
+    async def _clear_chatbot_database(self, user_id):
+        result = await self._clear_chatbot_history_in_db(user_id)
+        if result.modified_count > 0:
+            return "Chat history cleared successfully."
+        else:
+            return "No chat history found to clear."
+
+    async def set_chat_setting_antinsfw(self, chat_id, isboolean):
+        await self.antinsfw.update_one(
+            {"chat_id": chat_id},
+            {"$set": {"antinsfw": isboolean}},
+            upsert=True
+        )
+
+    async def chat_antinsfw(self, chat_id):
+        user_data = await self.antinsfw.find_one({"chat_id": chat_id})
+        if user_data:
+            return user_data.get("antinsfw", False)
+        return False
+
+
+db = Database(MONGO_URL)