From 3bdd7b629796a4634c2bd96400e01935955c14f8 Mon Sep 17 00:00:00 2001 From: Itz-fork Date: Thu, 28 Dec 2023 16:34:25 +0530 Subject: [PATCH] add admin features and update user controls - added ban / unban features - update database schema to keep track of files a user has downloaded / uploaded --- README.md | 8 ----- megadl/__init__.py | 2 +- megadl/helpers/database.py | 66 ++++++++++++++++++++++++++++++-------- megadl/helpers/mclient.py | 2 +- megadl/modules/admin.py | 45 ++++++++++++++++++++++++++ megadl/modules/mega_dl.py | 2 ++ megadl/modules/mega_up.py | 2 ++ 7 files changed, 103 insertions(+), 24 deletions(-) create mode 100644 megadl/modules/admin.py diff --git a/README.md b/README.md index 008eeae2..fd3ecf23 100644 --- a/README.md +++ b/README.md @@ -60,14 +60,6 @@ python3 -m megadl ### Config vars Please refer to [documentation](https://megabot.hirusha.codes/config-vars) -# Roadmap -- [x] Implement private mode -- [x] Implement DDL to Mega.nz uploader -- [ ] Better CLI output parser -- [x] Heroku support -- [x] Port [installer](https://github.com/Itz-fork/Mega.nz-Bot/blob/legacy/startup.sh) -- [x] Add better documentation - # Support [![Support Group](https://img.shields.io/badge/Support_Group-0a0a0a?style=for-the-badge&logo=telegram&logoColor=white)](https://t.me/Nexa_bots) \ No newline at end of file diff --git a/megadl/__init__.py b/megadl/__init__.py index 68c1c64d..376803f8 100644 --- a/megadl/__init__.py +++ b/megadl/__init__.py @@ -4,7 +4,7 @@ # Description: __init__.py # start msg -print("Mega.nz Bot is starting...") +print("Mega.nz Bot - Cypher is starting...") # loading config diff --git a/megadl/helpers/database.py b/megadl/helpers/database.py index cc42b05e..9913a534 100644 --- a/megadl/helpers/database.py +++ b/megadl/helpers/database.py @@ -5,56 +5,94 @@ from megadl.lib.aiomongo import AioMongo + class Users: def __init__(self) -> None: self.mongoc = AioMongo() - self.coll = self.mongoc["mega_nz"]["users"] + self.mdb = self.mongoc["mega_nz"] + self.coll_users = self.mdb["users"] + self.coll_banned = self.mdb["banned"] - # <<<<<<<<<< Telegram functions >>>>>>>>>> # + # <<<<<<<<<< Active user functions >>>>>>>>>> # async def add(self, user_id: int): added = await self.is_there(user_id) if not added: await self.mongoc.insert( - self.coll, {"_id": user_id, "email": "", "password": ""} + self.coll_users, + { + "_id": user_id, + "email": "", + "password": "", + "total_downloads": 0, + "total_uploads": 0, + }, + ) + + async def plus_fl_count( + self, user_id: int, downloads: int | None = None, uploads: int | None = None + ): + _, _, ctdl, cuptl = await self.is_there(user_id) + if downloads: + await self.mongoc.update( + self.coll_users, {"_id": user_id}, {"total_downloads": ctdl + downloads} + ) + elif uploads: + await self.mongoc.update( + self.coll_users, {"_id": user_id}, {"total_uploads": cuptl + uploads} ) async def delete(self, user_id: int): uid = {"_id": user_id} - added = await self.is_there(self.coll, user_id) + added = await self.is_there(self.coll_users, user_id) if added: - await self.mongoc.delete(self.coll, uid) + await self.mongoc.delete(self.coll_users, uid) async def is_there(self, user_id: int, use_acc: bool = False): uid = {"_id": user_id} - docu = await self.mongoc.find(self.coll, uid) + docu = await self.mongoc.find(self.coll_users, uid) if use_acc: email = docu["email"] password = docu["password"] - return [email, password] if not "" in {email, password} else None + tdl = docu["total_downloads"] + tupld = docu["total_uploads"] + return ( + [email, password, tdl, tupld] if not "" in {email, password} else None + ) else: - return docu + return docu + + # <<<<<<<<<< Banned user functions >>>>>>>>>> # + async def ban_user(self, user_id: int): + uid = {"_id": user_id} + await self.delete(user_id) + banned = await self.mongoc.find(self.coll_banned, user_id) + if not banned: + await self.mongoc.insert(self.coll_banned, {"_id": user_id}) + + async def unban_user(self, user_id: int): + await self.mongoc.delete(self.coll_banned, {"_id": user_id}) # <<<<<<<<<< Mega functions >>>>>>>>>> # async def mega_login(self, user_id: int, email: str, password: str): uid = {"_id": user_id} - logged = await self.mongoc.find(self.coll, uid) + logged = await self.mongoc.find(self.coll_users, uid) if logged: await self.mongoc.update( - self.coll, uid, {"email": email, "password": password} + self.coll_users, uid, {"email": email, "password": password} ) # this should never happen but better safe than sorry else: await self.mongoc.insert( - self.coll, {"_id": user_id, "email": email, "password": password} + self.coll_users, {"_id": user_id, "email": email, "password": password} ) async def mega_logout(self, user_id: int): uid = {"_id": user_id} - logged = await self.mongoc.find(self.coll, uid) + logged = await self.mongoc.find(self.coll_users, uid) if logged: - await self.mongoc.delete(self.coll, uid) + await self.mongoc.delete(self.coll_users, uid) async def how_many(self): - return (user["user_id"] async for user in self.coll.find({})) + return (user["user_id"] async for user in self.coll_users.find({})) diff --git a/megadl/helpers/mclient.py b/megadl/helpers/mclient.py index 25453648..577149c8 100644 --- a/megadl/helpers/mclient.py +++ b/megadl/helpers/mclient.py @@ -36,7 +36,7 @@ class MeganzClient(Client): Custom pyrogram client for Mega.nz-Bot """ - version = "v2-nightly" + version = "cypher-1.0" dl_loc = None tmp_loc = None database = Users() if os.getenv("MONGO_URI") else None diff --git a/megadl/modules/admin.py b/megadl/modules/admin.py new file mode 100644 index 00000000..3a7d9e3d --- /dev/null +++ b/megadl/modules/admin.py @@ -0,0 +1,45 @@ +# Copyright (c) 2023 Itz-fork +# Author: https://github.com/Itz-fork +# Project: https://github.com/Itz-fork/Mega.nz-Bot +# Description: Admin functions + +from pyrogram import filters +from pyrogram.types import Message + +from megadl import MeganzClient + + +@MeganzClient.on_message(filters.command("ban")) +@MeganzClient.handle_checks +async def admin_ban_user(client: MeganzClient, msg: Message): + sender = msg.from_user.id + if sender not in client.auth_users: + return await msg.reply("Banning users can only be done by admins!") + buid = None + try: + buid = int(msg.text.split(None, 1)[1]) + except: + pass + if not buid or not buid.isnumeric(): + return await msg.reply("Provide a user id to ban \n\nEx: `/ban 12345`") + + await client.database.ban_user(buid) + await msg.reply(f"Banned user `{buid}`") + + +@MeganzClient.on_message(filters.command("unban")) +@MeganzClient.handle_checks +async def admin_unban_user(client: MeganzClient, msg: Message): + sender = msg.from_user.id + if sender not in client.auth_users: + return await msg.reply("Unbanning users can only be done by admins!") + buid = None + try: + buid = int(msg.text.split(None, 1)[1]) + except: + pass + if not buid or not buid.isnumeric(): + return await msg.reply("Provide a user id to unban \n\nEx: `/ban 12345`") + + await client.database.unban_user(buid) + await msg.reply(f"Unbanned user `{buid}`") diff --git a/megadl/modules/mega_dl.py b/megadl/modules/mega_dl.py index e630e76c..ec44afdb 100644 --- a/megadl/modules/mega_dl.py +++ b/megadl/modules/mega_dl.py @@ -102,6 +102,8 @@ async def dl_from_cb(client: MeganzClient, query: CallbackQuery): `{e}` """ ) + # update download count + await client.database.plus_fl_count(qcid, downloads=len(f_list)) # Send file(s) to the user await resp.edit("Trying to upload now 📤...") await client.send_files(f_list, qcid, resp.id) diff --git a/megadl/modules/mega_up.py b/megadl/modules/mega_up.py index e927f20d..381be577 100644 --- a/megadl/modules/mega_up.py +++ b/megadl/modules/mega_up.py @@ -68,6 +68,8 @@ async def to_up_cb(client: MeganzClient, query: CallbackQuery): await client.edit_message_text( qcid, qmid, "Trying to download the file 📥", reply_markup=None ) + # update upload count + await client.database.plus_fl_count(qcid, uploads=1) # Download files accordingly dl_path = None