From ad166c397a57ca0bdf38a82a2eb72581b0b62160 Mon Sep 17 00:00:00 2001 From: Itz-fork Date: Sun, 31 Dec 2023 17:25:44 +0530 Subject: [PATCH] new method and some other changes - added user account info method - store env in a variable instead of calling copy method on each request - paste folder struct to nekobin to avoid message too long errors --- megadl/helpers/cypher.py | 1 + megadl/helpers/sysfncs.py | 38 +++++++++++++++++++++++++++++++++----- megadl/lib/megatools.py | 23 +++++++++++++++++++---- megadl/modules/mega_dl.py | 19 +++++++++++++------ megadl/modules/mega_up.py | 31 +++++++++++++++++++++++++++++++ 5 files changed, 97 insertions(+), 15 deletions(-) diff --git a/megadl/helpers/cypher.py b/megadl/helpers/cypher.py index fe2082a3..fc42d4a1 100644 --- a/megadl/helpers/cypher.py +++ b/megadl/helpers/cypher.py @@ -77,6 +77,7 @@ def __init__(self): # Initializing mongodb print("> Initializing database") self.glob_tmp = {} + self.environs = os.environ.copy() self.cipher = None print("> Updating privacy settings") diff --git a/megadl/helpers/sysfncs.py b/megadl/helpers/sysfncs.py index 8495ec3b..58360735 100644 --- a/megadl/helpers/sysfncs.py +++ b/megadl/helpers/sysfncs.py @@ -3,8 +3,8 @@ # Project: https://github.com/Itz-fork/Mega.nz-Bot # Description: Shell, loops and other sys functions +import asyncio -from asyncio import to_thread from functools import partial from subprocess import Popen, PIPE from psutil import Process as PsuPrc @@ -15,9 +15,11 @@ async def run_partial(func, *args, **kwargs): """ Run synchronous functions in a non-blocking coroutine and return the result - Arguments: - + - Arguments: - func: function - Function to execute + + - Note: + Used for CPU-bound tasks or blocking IO that can't be handled efficiently with asyncio's non-blocking IO. """ if iscoroutinefunction(func): return await func(*args, **kwargs) @@ -26,15 +28,41 @@ async def run_partial(func, *args, **kwargs): return await loop.run_in_executor(None, partial(func, *args, **kwargs)) -def run_on_shell(cmd): +def run_on_shell(cmd: str): """ Run shell commands and get it's output + + - Arguments: + - cmd: str - Command to execute """ run = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True) shout = run.stdout.read()[:-1].decode("utf-8") return shout +async def with_sub_shell(cmd: str, **kwargs): + """ + Run shell commands and get it's output + + - Arguments: + - cmd: str - Command to execute + + - Note: + Creates a subprocess using asyncio.create_subprocess_shell + Not suitable for heavy IO bound tasks (use run_on_shell with run_partial for that) + """ + process = await asyncio.create_subprocess_shell( + cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, **kwargs + ) + + stdout, stderr = await process.communicate() + + if stdout: + return stdout[:-1].decode("utf-8") + if stderr: + return stderr[:-1].decode("utf-8") + + async def kill_family(pid: int): """ Murder a process and it's whole family using pid @@ -43,4 +71,4 @@ async def kill_family(pid: int): for child in running.children(recursive=True): child.kill() running.kill() - await to_thread(running.wait) + await asyncio.to_thread(running.wait) diff --git a/megadl/lib/megatools.py b/megadl/lib/megatools.py index 648377d7..36b6e3af 100644 --- a/megadl/lib/megatools.py +++ b/megadl/lib/megatools.py @@ -11,7 +11,7 @@ from humans import human_bytes from aiohttp import ClientSession from megadl.helpers.files import listfiles -from megadl.helpers.sysfncs import run_partial, run_on_shell +from megadl.helpers.sysfncs import run_partial, run_on_shell, with_sub_shell from megadl.helpers.crypt import ( base64_to_a32, decrypt_attr, @@ -36,6 +36,10 @@ class MegaRegexs: file_key = re.compile(r"(?<=#)(.*)") old_f_ik = re.compile(r"(?<=!)(.*)") + user_total = re.compile(r"Total: (.*)") + user_used = re.compile(r"Used: (.*)") + user_free = re.compile(r"Free: (.*)") + Regexes = MegaRegexs() @@ -156,6 +160,17 @@ async def upload( ) return ulink + async def user_fs(self): + """ """ + _sh_dta = await with_sub_shell(f"megadf -h {self.config}") + if _sh_dta: + _total = Regexes.user_total.search(_sh_dta).group(1) + _used = Regexes.user_used.search(_sh_dta).group(1) + _free = Regexes.user_free.search(_sh_dta).group(1) + return _total, _used, _free + else: + raise LoginError + @staticmethod async def get_info(url: str) -> list[str]: """ @@ -293,7 +308,7 @@ async def __shellExec( cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, - env=os.environ.copy(), + env=self.client.environs, ) self.client.mega_running[user_id] = run.pid @@ -312,8 +327,8 @@ async def handle_stdout(out): await self.client.edit_message_text( chat_id, msg_id, f"**Process info:** \n`{out}`", **kwargs ) - except: - pass + except Exception as e: + print(e) async def handle_stderr(err): if run.returncode is None: diff --git a/megadl/modules/mega_dl.py b/megadl/modules/mega_dl.py index 6cf5a452..fbbaed76 100644 --- a/megadl/modules/mega_dl.py +++ b/megadl/modules/mega_dl.py @@ -6,6 +6,7 @@ import re from os import path, makedirs +from aiohttp import ClientSession from pyrogram import filters from pyrogram.types import ( @@ -127,9 +128,15 @@ async def info_from_cb(client: CypherClient, query: CallbackQuery): ) else: - await query.edit_message_text( - f""" -`{retrieved}` -""", - reply_markup=None, - ) + 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 d61cee90..2bd000bb 100644 --- a/megadl/modules/mega_up.py +++ b/megadl/modules/mega_up.py @@ -102,3 +102,34 @@ 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}` +""")