proxy support and bug fixes

- added proxy support for megadl
- added user account info command
- renamed database to CypherDB
- fixed database resetting everytime a user starts the bot
- fixed folder info error
- moved additional functions to bonus.py
- updated message strings

megadl/helpers/cypher.py

@@ -16,7 +16,7 @@
 from pyrogram import Client, errors
 from pyrogram.handlers import MessageHandler
 
-from .database import Users
+from .database import CypherDB
 from .files import send_as_guessed, fs_cleanup, splitit, listfiles
 
 
@@ -39,7 +39,7 @@ class MeganzClient(Client):
     version = "cypher-1.0"
     dl_loc = None
     tmp_loc = None
-    database = Users() if os.getenv("MONGO_URI") else None
+    database = CypherDB() if os.getenv("MONGO_URI") else None
 
     def __init__(self):
         # set DOWNLOAD_LOCATION variable
@@ -146,7 +146,11 @@ async def cy_run(client: Client, msg: Message):
 
                 # Check auth users
                 if self.database:
-                    await self.database.add(uid)
+                    status = await self.database.add(uid)
+                    if status["banned"]:
+                        return await msg.reply(
+                            f"**You're banned from using this bot 😬** \n\n**Reason:** `{status['reason']}`"
+                        )
 
                 if "*" in self.auth_users:
                     can_use = True

megadl/helpers/database.py

@@ -6,7 +6,7 @@
 from megadl.lib.aiomongo import AioMongo
 
 
-class Users:
+class CypherDB:
     def __init__(self) -> None:
         self.mongoc = AioMongo()
         self.coll_users = self.mongoc["mega_nz"]["users"]
@@ -24,9 +24,12 @@ async def add(self, user_id: int):
                 "status": {"banned": False, "reason": ""},
                 "total_downloads": 0,
                 "total_uploads": 0,
+                "proxy": "",
             },
+            no_modify=True,
             upsert=True,
         )
+        return await self.mongoc.find_async(self.coll_users, {"_id": user_id}, {"status": 1})
 
     async def plus_fl_count(
         self, user_id: int, downloads: int | None = None, uploads: int | None = None
@@ -56,6 +59,7 @@ async def is_there(self, user_id: int, use_acc: bool = False):
                 "status": {"banned": False, "reason": ""},
                 "total_downloads": 0,
                 "total_uploads": 0,
+                "proxy": ""
             }
         """
         uid = {"_id": user_id}
@@ -64,7 +68,7 @@ async def is_there(self, user_id: int, use_acc: bool = False):
             return None
         docu.pop("_id")
         if use_acc:
-            return docu if not "" in docu.values() else None
+            return docu if not "" in {docu["email"], docu["password"]} else None
         else:
             return docu
 
@@ -86,6 +90,7 @@ async def unban_user(self, user_id: int):
     # <<<<<<<<<< Mega functions >>>>>>>>>> #
 
     async def mega_login(self, user_id: int, email: str, password: str):
+        print(user_id, email, password)
         await self.mongoc.update_async(
             self.coll_users,
             {"_id": user_id},
@@ -98,3 +103,17 @@ async def mega_logout(self, user_id: int):
 
     async def how_many(self):
         return (user["user_id"] async for user in self.coll_users.find({}))
+
+    # <<<<<<<<<< Proxy functions >>>>>>>>>> #
+    async def update_proxy(self, user_id: int, proxy: str):
+        await self.mongoc.update_async(
+            self.coll_users,
+            {"_id": user_id},
+            {"proxy": proxy},
+            upsert=True,
+        )
+
+    async def get_proxy(self, user_id: int):
+        return await self.mongoc.find_async(
+            self.coll_users, {"_id": user_id}, {"proxy": 1}
+        )

megadl/lib/aiomongo.py

@@ -26,7 +26,13 @@ def __init__(
     ) -> None:
         self.atlas_host = getenv("MONGO_URI")
         super().__init__(
-            self.atlas_host, port, document_class, tz_aware, connect, type_registry, **kwargs
+            self.atlas_host,
+            port,
+            document_class,
+            tz_aware,
+            connect,
+            type_registry,
+            **kwargs
         )
 
     async def insert_async(self, coll: Collection, query: dict, *args, **kwargs):
@@ -41,16 +47,43 @@ async def find_async(self, coll: Collection, query: dict, *args, **kwargs):
         """
         return await run_partial(coll.find_one, query, *args, **kwargs)
 
-    async def update_async(self, coll: Collection, query: dict, value: dict, *args, **kwargs):
+    async def update_async(
+        self,
+        coll: Collection,
+        query: dict,
+        value: dict,
+        no_modify: bool = False,
+        *args,
+        **kwargs
+    ):
         """
         Perform `update_one` operation on the given collection
         """
-        return await run_partial(
-            coll.update_one, query, {"$set": value}, *args, **kwargs
-        )
+        if no_modify:
+            return await run_partial(
+                coll.update_one, query, {"$setOnInsert": value}, *args, **kwargs
+            )
+        else:
+            return await run_partial(
+                coll.update_one, query, {"$set": value}, *args, **kwargs
+            )
 
     async def delete_async(self, coll: Collection, query: dict, *args, **kwargs):
         """
         Perform `delete_one` operation on the given collection
         """
         return await run_partial(coll.delete_one, query, *args, **kwargs)
+
+    async def count_documents_async(
+        self, coll: Collection, query: dict, *args, **kwargs
+    ):
+        """
+        Perform `count_documents` operation on the given collection
+        """
+        return await run_partial(coll.count_documents, query, *args, **kwargs)
+
+    async def find_many_async(self, coll: Collection, query: dict, *args, **kwargs):
+        """
+        Perform `find` operation on the given collection
+        """
+        return await run_partial(coll.find, query, *args, **kwargs)

megadl/lib/megatools.py

@@ -40,6 +40,8 @@ class MegaRegexs:
     user_used = re.compile(r"Used: (.*)")
     user_free = re.compile(r"Free: (.*)")
 
+    proxy_regex = re.compile(r'^(http|https|socks5|socks5h)://([a-zA-Z0-9\-\.]+):(\d+)$')
+
 
 Regexes = MegaRegexs()
 
@@ -279,7 +281,7 @@ async def prepare_string(nodes, shared_key, depth=0):
                             )
                             attrs = decrypt_attr(base64_url_decode(node["a"]), k)
                             file_name = attrs["n"]
-                            to_return += f"{' ' * depth}├─ {file_name} ({human_bytes(file_size)})\n"
+                            to_return += f"{' ' * depth}├── {file_name} ({human_bytes(file_size)})\n"
                         elif node["t"] == 1:  # folder
                             k = key
                             attrs = decrypt_attr(base64_url_decode(node["a"]), k)
@@ -304,6 +306,7 @@ async def prepare_string(nodes, shared_key, depth=0):
     async def __shellExec(
         self, cmd: str, user_id: int, chat_id: int = None, msg_id: int = None, **kwargs
     ):
+        print(cmd)
         run = await asyncio.create_subprocess_shell(
             cmd,
             stdout=asyncio.subprocess.PIPE,

megadl/modules/bonus.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2023 Itz-fork
+# Author: https://github.com/Itz-fork
+# Project: https://github.com/Itz-fork/Mega.nz-Bot
+# Description: Bonus functions such as file/folder info, proxy setting and account status
+
+from pyrogram import filters
+from aiohttp import ClientSession
+from pyrogram.types import (
+    Message,
+    CallbackQuery,
+    InlineKeyboardButton,
+    InlineKeyboardMarkup,
+)
+
+from megadl import CypherClient
+from megadl.lib.megatools import Regexes, MegaTools
+
+
+@CypherClient.on_callback_query(filters.regex(r"info_mg?.+"))
+@CypherClient.run_checks
+async def info_from_cb(client: CypherClient, query: CallbackQuery):
+    url = client.glob_tmp.get(query.from_user.id)[0]
+    await query.edit_message_text("`Getting info ℹ️...`")
+    retrieved = await MegaTools.get_info(url)
+
+    if isinstance(retrieved, list):
+        await query.edit_message_text(
+            f"""
+》 **File Details**
+
+**📛 Name:** `{retrieved[0]}`
+**🗂 Size:** `{retrieved[1]}`
+**📎 URL:** `{url}`
+""",
+            reply_markup=None,
+        )
+
+    else:
+        async with ClientSession() as nekoc:
+            resp = await nekoc.post(
+                "https://nekobin.com/api/documents", json={"content": retrieved}
+            )
+            if resp.status == 201:
+                nekourl = f"https://nekobin.com/{(await resp.json())['result']['key']}"
+                await query.edit_message_text(
+                    f"**👀 View is the generated [folder info]({nekourl})**",
+                    reply_markup=InlineKeyboardMarkup(
+                        [[InlineKeyboardButton("Visit 🔗", url=nekourl)]]
+                    ),
+                )
+            else:
+                await query.edit_message_text(
+                    "`Failed to generate folder info 😕, Please try again later`"
+                )
+
+
+@CypherClient.on_message(filters.command("acc"))
+@CypherClient.run_checks
+async def acc(_: CypherClient, msg: Message):
+    # Check if the user exists in database
+    usr = msg.from_user.id
+    udoc = await _.database.is_there(usr, True)
+    if not udoc:
+        return await msg.reply(
+            "`You must be logged in first to check account status 😑`"
+        )
+
+    # Get user data
+    email = _.cipher.decrypt(udoc["email"]).decode()
+    password = _.cipher.decrypt(udoc["password"]).decode()
+    conf = f"--username {email} --password {password}"
+    cli = MegaTools(_, conf)
+    total, used, free = await cli.user_fs()
+
+    return await msg.reply(
+        f"""
+**~ Your User Account Info ~**
+
+✦ **Email:** `{email}`
+✦ **Password:** `{password}`
+✦ **Storage,**
+    ⤷ **Total:** `{total}`
+    ⤷ **Used:** `{used}`
+    ⤷ **Free:** `{free}`
+"""
+    )
+
+
+@CypherClient.on_message(filters.command("proxy"))
+@CypherClient.run_checks
+async def set_user_proxy(client: CypherClient, msg: Message):
+    prxy = None
+    try:
+        prxy = msg.text.split(None, 1)[1]
+    except:
+        pass
+    if not prxy or not Regexes.proxy_regex.match(prxy):
+        return await msg.reply(
+            "Provide a proxy to set \n\nEx: `/proxy socks5h://localhost:9050`"
+        )
+
+    await client.database.update_proxy(msg.from_user.id, prxy)
+    await msg.reply(
+        f"Proxy set to: `{prxy}` \n\n**Note:**\nIf your download seems to stuck, it's likely due to proxy not working properly. Free proxies on internet do not work!"
+    )

megadl/modules/generals.py

@@ -37,11 +37,18 @@ async def help_msg(_: CypherClient, msg: Message):
 
 **✘ How to download from mega link?**
   ⤷ It's very easy. Just send the link you want to download and I'll download it for you 😉.
-  For private content you need to login first then send path to the file or folder you want to download starting with `/Root/<path-to-dl>`.
+  ⤷ For private content you need to login first then send path to the file or folder you want to download starting with `/Root/<path-to-dl>`.
 
 **✘ How to upload files to Mega.nz?**
   ⤷ Just send me the files and I'll ask you whether you want to upload it or not. Same goes for direct download links 😎
 
+**✘ How to setup proxy?**
+  ⤷ Send /proxy command alongside the proxy 📡 (Ex: `/proxy https://example.com:8080`)
+    Please note that the free proxies you see on internet are **not working**
+  
+** ✘ How to get my account details?**
+  ⤷ Send /acc command and I'll send you your account details 🫣
+
 
 **Made with ❤️ by @NexaBotsUpdates**
       """

megadl/modules/mega_dl.py

@@ -6,7 +6,6 @@
 
 import re
 from os import path, makedirs
-from aiohttp import ClientSession
 
 from pyrogram import filters
 from pyrogram.types import (
@@ -64,7 +63,10 @@ async def dl_from_cb(client: CypherClient, query: CallbackQuery):
                 "`You must be logged in first to download this file 😑`"
             )
         if udoc:
-            conf = f"--username {client.cipher.decrypt(udoc['email']).decode()} --password {client.cipher.decrypt(udoc['password']).decode()}"
+            email = client.cipher.decrypt(udoc["email"]).decode()
+            password = client.cipher.decrypt(udoc["password"]).decode()
+            proxy = f"--proxy {udoc['proxy']}" if udoc["proxy"] else ""
+            conf = f"--username {email} --password {password} {proxy}"
 
     # Create unique download folder
     if not path.isdir(dlid):
@@ -107,36 +109,3 @@ async def dl_from_cb(client: CypherClient, query: CallbackQuery):
     )
     await client.full_cleanup(dlid, qusr)
     await resp.delete()
-
-
-@CypherClient.on_callback_query(filters.regex(r"info_mg?.+"))
-@CypherClient.run_checks
-async def info_from_cb(client: CypherClient, query: CallbackQuery):
-    url = client.glob_tmp.get(query.from_user.id)[0]
-    retrieved = await MegaTools.get_info(url)
-
-    if isinstance(retrieved, list):
-        await query.edit_message_text(
-            f"""
-》 **File Details**
-
-**📛 Name:** `{retrieved[0]}`
-**🗂 Size:** `{retrieved[1]}`
-**📎 URL:** `{url}`
-""",
-            reply_markup=None,
-        )
-
-    else:
-        async with ClientSession() as nekoc:
-            resp = await nekoc.post(
-                "https://nekobin.com/api/documents", json={"content": retrieved}
-            )
-            if resp.status == 200:
-                nekourl = f"https://nekobin.com/{(await resp.json())['result']['key']}"
-                await query.edit_message_text(
-                    f"Check folder info here: `{nekourl}`",
-                    reply_markup=InlineKeyboardMarkup(
-                        [[InlineKeyboardButton("Visit 🔗", url=nekourl)]]
-                    ),
-                )