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

@@ -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)

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

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({}))

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

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}`")

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)

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