diff --git a/megadl/helpers/cypher.py b/megadl/helpers/cypher.py index fc42d4a1..606c86b6 100644 --- a/megadl/helpers/cypher.py +++ b/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 diff --git a/megadl/helpers/database.py b/megadl/helpers/database.py index 2315b5be..c60719cb 100644 --- a/megadl/helpers/database.py +++ b/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} + ) diff --git a/megadl/lib/aiomongo.py b/megadl/lib/aiomongo.py index d3ffd526..810986c2 100644 --- a/megadl/lib/aiomongo.py +++ b/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) diff --git a/megadl/lib/megatools.py b/megadl/lib/megatools.py index 36b6e3af..5b6945d3 100644 --- a/megadl/lib/megatools.py +++ b/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, diff --git a/megadl/modules/bonus.py b/megadl/modules/bonus.py new file mode 100644 index 00000000..46c56a1d --- /dev/null +++ b/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!" + ) diff --git a/megadl/modules/generals.py b/megadl/modules/generals.py index a9e66a8f..cb24beb8 100644 --- a/megadl/modules/generals.py +++ b/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/`. + โคท For private content you need to login first then send path to the file or folder you want to download starting with `/Root/`. **โœ˜ 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** """ diff --git a/megadl/modules/mega_dl.py b/megadl/modules/mega_dl.py index fbbaed76..95957db6 100644 --- a/megadl/modules/mega_dl.py +++ b/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)]] - ), - ) diff --git a/megadl/modules/mega_up.py b/megadl/modules/mega_up.py index 2bd000bb..d61cee90 100644 --- a/megadl/modules/mega_up.py +++ b/megadl/modules/mega_up.py @@ -102,34 +102,3 @@ async def to_up_cb(client: CypherClient, query: CallbackQuery): ), ) await client.full_cleanup(dl_path, qusr) - - -# CypherClient on "/acc" command -@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}` -""")