diff --git a/sukuinote/plugins/pyexec.py b/sukuinote/plugins/pyexec.py index d2375dd..aa0ee36 100644 --- a/sukuinote/plugins/pyexec.py +++ b/sukuinote/plugins/pyexec.py @@ -9,6 +9,8 @@ from io import StringIO, BytesIO from pyrogram import Client, filters from .. import config, help_dict, log_errors, slave, apps, session, public_log_errors +exec_tasks = dict() + PYEXEC_REGEX = '^(?:' + '|'.join(map(re.escape, config['config']['prefixes'])) + r')exec\s+([\s\S]+)$' @Client.on_message(~filters.forwarded & ~filters.sticker & ~filters.via_bot & ~filters.edited & filters.me & filters.regex(PYEXEC_REGEX)) @log_errors @@ -39,8 +41,15 @@ async def pyexec(client, message): if ex.msg != "'return' with value in async generator": raise exx = _gf(obody) - reply = await message.reply_text('Executing...') - async_obj = exx(client, client, message, message, reply, message.reply_to_message, message.reply_to_message, UniqueExecReturnIdentifier) + reply = await message.reply_text(f'Executing {hash(UniqueExecReturnIdentifier)}...') + oasync_obj = exx(client, client, message, message, reply, message.reply_to_message, message.reply_to_message, UniqueExecReturnIdentifier) + if inspect.isasyncgen(oasync_obj): + async def async_obj(): + return [i async for i in oasync_obj] + else: + async def async_obj(): + to_return = [await oasync_obj] + return [] if to_return == [UniqueExecReturnIdentifier] else to_return stdout = sys.stdout stderr = sys.stderr wrapped_stdout = StringIO() @@ -48,15 +57,19 @@ async def pyexec(client, message): try: sys.stdout = wrapped_stdout sys.stderr = wrapped_stderr - if inspect.isasyncgen(async_obj): - returned = [i async for i in async_obj] - else: - returned = [await async_obj] - if returned == [UniqueExecReturnIdentifier]: - returned = [] + task = asyncio.create_task(async_obj()) + exec_tasks[hash(UniqueExecReturnIdentifier)] = task + returned = await async_obj + except asyncio.CancelledError: + sys.stdout = stdout + sys.stderr = stderr + exec_tasks.pop(hash(UniqueExecReturnIdentifier), None) + await reply.edit_text('Cancelled') + return finally: sys.stdout = stdout sys.stderr = stderr + exec_tasks.pop(hash(UniqueExecReturnIdentifier), None) wrapped_stderr.seek(0) wrapped_stdout.seek(0) output = '' @@ -83,4 +96,36 @@ async def pyexec(client, message): else: await reply.edit_text(output) -help_dict['exec'] = ('Exec', '{prefix}exec <python code> - Executes python code') +@Client.on_message(~filters.forwarded & ~filters.sticker & ~filters.via_bot & ~filters.edited & filters.me & filters.command(['listexecs', 'listexec', 'lexec'], prefixes=config['config']['prefixes'])) +@log_errors +@public_log_errors +async def listexec(client, message): + text = '\n'.join(map(str, exec_tasks)) + if len(text) > 4096: + f = BytesIO(text.encode('utf-8')) + f.name = 'exectasks.txt' + await message.reply_document(f) + else: + text = '\n'.join(map(lambda i: f'{i}', exec_tasks)) + await message.reply_text(text) + +@Client.on_message(~filters.forwarded & ~filters.sticker & ~filters.via_bot & ~filters.edited & filters.me & filters.command(['cancelexec', 'cexec'], prefixes=config['config']['prefixes'])) +@log_errors +@public_log_errors +async def cancelexec(client, message): + try: + task = exec_tasks.get(int(message[1])) + except IndexError: + return + if not task: + await message.reply_text('Task does not exist') + return + task.cancel() + +help_dict['exec'] = ('Exec', '''{prefix}exec <python code> - Executes python code + +{prefix}listexecs - List exec tasks +Aliases: {prefix}listexec, {prefix}lexec + +{prefix}cancelexec <task id> - Cancel exec task +Aliases: {prefix}cexec''')