refactoring, privacy and database updates

- database: update database schema to only use "users" collection for everything including bans
- privacy: support for private, public and global mode with admins
- admin: added info command

some other minor changes

megadl/__init__.py

@@ -13,4 +13,6 @@
 print("> Loading config")
 load_dotenv()
 
-from .helpers.mclient import MeganzClient
+# client
+from .helpers.mclient import MeganzClient
+MegaCypher: "MeganzClient" = MeganzClient()

megadl/__main__.py

@@ -6,11 +6,10 @@
 
 from pyrogram import idle
 
-from . import MeganzClient
+from . import MegaCypher
 
 # Run the bot
 if __name__ == "__main__":
     # Custom pyrogram client
-    meganzbot = MeganzClient()
-    meganzbot.start()
+    MegaCypher.start()
     idle()

megadl/helpers/database.py

@@ -9,21 +9,20 @@
 class Users:
     def __init__(self) -> None:
         self.mongoc = AioMongo()
-        self.mdb = self.mongoc["mega_nz"]
-        self.coll_users = self.mdb["users"]
-        self.coll_banned = self.mdb["banned"]
+        self.coll_users = self.mongoc["mega_nz"]["users"]
 
     # <<<<<<<<<< Active user functions >>>>>>>>>> #
 
     async def add(self, user_id: int):
         added = await self.is_there(user_id)
         if not added:
-            await self.mongoc.insert(
+            await self.mongoc.insert_async(
                 self.coll_users,
                 {
                     "_id": user_id,
                     "email": "",
                     "password": "",
+                    "status": {"banned": False, "reason": ""},
                     "total_downloads": 0,
                     "total_uploads": 0,
                 },
@@ -32,67 +31,88 @@ async def add(self, user_id: int):
     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)
+        dl_up = await self.is_there(user_id)
         if downloads:
-            await self.mongoc.update(
-                self.coll_users, {"_id": user_id}, {"total_downloads": ctdl + downloads}
+            await self.mongoc.update_async(
+                self.coll_users,
+                {"_id": user_id},
+                {"total_downloads": dl_up["total_downloads"] + downloads},
             )
         elif uploads:
-            await self.mongoc.update(
-                self.coll_users, {"_id": user_id}, {"total_uploads": cuptl + uploads}
+            await self.mongoc.update_async(
+                self.coll_users,
+                {"_id": user_id},
+                {"total_uploads": dl_up["total_uploads"] + uploads},
             )
 
     async def delete(self, user_id: int):
         uid = {"_id": user_id}
         added = await self.is_there(self.coll_users, user_id)
         if added:
-            await self.mongoc.delete(self.coll_users, uid)
+            await self.mongoc.delete_async(self.coll_users, uid)
 
     async def is_there(self, user_id: int, use_acc: bool = False):
+        """
+        Returns:
+            {
+                "email": "",
+                "password": "",
+                "status": {"banned": False, "reason": ""},
+                "total_downloads": 0,
+                "total_uploads": 0,
+            }
+        """
         uid = {"_id": user_id}
-        docu = await self.mongoc.find(self.coll_users, uid)
+        docu = await self.mongoc.find_async(self.coll_users, uid)
+        if not docu:
+            return None
+        docu.pop("_id")
         if use_acc:
-            email = docu["email"]
-            password = docu["password"]
-            tdl = docu["total_downloads"]
-            tupld = docu["total_uploads"]
-            return (
-                [email, password, tdl, tupld] if not "" in {email, password} else None
-            )
+            return docu if not "" in docu.values() else None
         else:
             return docu
 
     # <<<<<<<<<< Banned user functions >>>>>>>>>> #
-    async def ban_user(self, user_id: int):
+    async def ban_user(self, user_id: int, reason: str):
         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})
+        to_ban = await self.mongoc.find_async(self.coll_users, uid)
+        if to_ban:
+            await self.mongoc.update_async(
+                self.coll_users,
+                {"_id": user_id},
+                {"status": {"banned": True, "reason": reason}},
+            )
 
     async def unban_user(self, user_id: int):
-        await self.mongoc.delete(self.coll_banned, {"_id": user_id})
+        uid = {"_id": user_id}
+        to_ban = await self.mongoc.find_async(self.coll_users, uid)
+        if to_ban:
+            await self.mongoc.update_async(
+                self.coll_users,
+                {"_id": user_id},
+                {"status": {"banned": False, "reason": "Got unbanned"}},
+            )
 
     # <<<<<<<<<< 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_users, uid)
+        logged = await self.mongoc.find_async(self.coll_users, uid)
         if logged:
-            await self.mongoc.update(
+            await self.mongoc.update_async(
                 self.coll_users, uid, {"email": email, "password": password}
             )
         # this should never happen but better safe than sorry
         else:
-            await self.mongoc.insert(
+            await self.mongoc.insert_async(
                 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_users, uid)
+        logged = await self.mongoc.find_async(self.coll_users, uid)
         if logged:
-            await self.mongoc.delete(self.coll_users, uid)
+            await self.mongoc.delete_async(self.coll_users, uid)
 
     async def how_many(self):
         return (user["user_id"] async for user in self.coll_users.find({}))

megadl/helpers/files.py

@@ -73,7 +73,7 @@ async def send_as_guessed(client, file, chat_id, mid):
             # Get duration of video in seconds
             _sh = await run_partial(
                 run_on_shell,
-                f"ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 {file}",
+                f"ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 '{file}'",
             )
             vid_dur = int(float(_sh))
             # Generate thumbnail for the video
@@ -83,7 +83,7 @@ async def send_as_guessed(client, file, chat_id, mid):
                 makedirs(_tmpth)
             _sh = await run_partial(
                 run_on_shell,
-                f"ffmpeg -i {file} -ss {timedelta(seconds=int(vid_dur/10))} -vframes 1 {_thumb}",
+                f"ffmpeg -y -ss {timedelta(seconds=int(vid_dur/10))} -i '{file}' -vframes 1 '{_thumb}'",
             )
             await client.send_video(
                 chat_id,

megadl/helpers/mclient.py

@@ -40,11 +40,6 @@ class MeganzClient(Client):
     dl_loc = None
     tmp_loc = None
     database = Users() if os.getenv("MONGO_URI") else None
-    auth_users = (
-        set(map(int, os.getenv("AUTH_USERS").split()))
-        if os.getenv("AUTH_USERS") and os.getenv("AUTH_USERS") != "*"
-        else "*"
-    )
 
     def __init__(self):
         # set DOWNLOAD_LOCATION variable
@@ -69,7 +64,7 @@ def __init__(self):
         # Initializing pyrogram
         print("> Initializing client")
         super().__init__(
-            "MegaBot",
+            "MegaBotCypher",
             bot_token=os.getenv("BOT_TOKEN"),
             api_id=os.getenv("APP_ID"),
             api_hash=os.getenv("API_HASH"),
@@ -81,6 +76,23 @@ def __init__(self):
         print("> Initializing database")
         self.glob_tmp = {}
         self.cipher = None
+
+        print("> Updating privacy settings")
+        _auths = os.getenv("AUTH_USERS")
+        self.auth_users = (
+            set()
+            if not isinstance(_auths, (list, str))
+            else set(map(int, os.getenv("AUTH_USERS").split()))
+            if not _auths.startswith("*")
+            else {"*"}
+            if len(_auths.split("|")) > 2
+            else set(map(int, _auths.split("|")[1].split())).union({"*"})
+        )
+
+        self.use_logs = (
+            {"dl_from", "up_to"} if os.getenv("USE_LOGS") in {"True", "true"} else set()
+        )
+        self.log_chat = int(os.getenv("LOG_CHAT")) if os.getenv("LOG_CHAT") else None
         self.is_public = True if self.database else False
 
         if self.is_public:
@@ -111,22 +123,21 @@ def __init__(self):
 
         print("--------------------")
 
-    @classmethod
-    def handle_checks(self, func: Callable) -> Callable:
+    def run_checks(self, func) -> Callable:
         """
         Decorator to run middleware
         """
 
-        async def fn_run(client: Client, msg: Message):
+        async def cy_run(client: Client, msg: Message):
             can_use = False
             uid = msg.from_user.id
             try:
+                # Check auth users
                 if self.database:
                     await self.database.add(uid)
 
-                if self.auth_users == "*":
+                if "*" in self.auth_users:
                     can_use = True
-
                 else:
                     can_use = uid in self.auth_users
 
@@ -136,6 +147,14 @@ async def fn_run(client: Client, msg: Message):
                     )
                     return msg.stop_propagation()
 
+                # send logs if needed
+                if func.__name__ in self.use_logs:
+                    if self.log_chat:
+                        _frwded = await msg.forward(chat_id=self.log_chat)
+                        await _frwded.reply(
+                            f"**#UPLOAD_LOG** \n\n**From:** `{uid}` \n**Get history:** `/info {uid}`"
+                        )
+
                 return await func(client, msg)
             # Floodwait handling
             except errors.FloodWait as e:
@@ -148,7 +167,7 @@ async def fn_run(client: Client, msg: Message):
             except Exception as e:
                 logging.warning(_emsg.format(self.version, func.__module__, e))
 
-        return fn_run
+        return cy_run
 
     async def ask(self, chat_id: int, text: str, *args, **kwargs):
         await self.send_message(chat_id, text, *args, **kwargs)

megadl/lib/aiomongo.py

@@ -29,27 +29,27 @@ def __init__(
             self.atlas_host, port, document_class, tz_aware, connect, type_registry, **kwargs
         )
 
-    async def insert(self, coll: Collection, query: dict, *args, **kwargs):
+    async def insert_async(self, coll: Collection, query: dict, *args, **kwargs):
         """
         Perform `insert_one` operation on the given collection
         """
         return await run_partial(coll.insert_one, query, *args, **kwargs)
 
-    async def find(self, coll: Collection, query: dict, *args, **kwargs):
+    async def find_async(self, coll: Collection, query: dict, *args, **kwargs):
         """
         Perform `find_one` operation on the given collection
         """
         return await run_partial(coll.find_one, query, *args, **kwargs)
 
-    async def update(self, coll: Collection, query: dict, value: dict, *args, **kwargs):
+    async def update_async(self, coll: Collection, query: dict, value: dict, *args, **kwargs):
         """
         Perform `update_one` operation on the given collection
         """
         return await run_partial(
             coll.update_one, query, {"$set": value}, *args, **kwargs
         )
 
-    async def delete(self, coll: Collection, query: dict, *args, **kwargs):
+    async def delete_async(self, coll: Collection, query: dict, *args, **kwargs):
         """
         Perform `delete_one` operation on the given collection
         """