Add exec tasks

This commit is contained in:
blank X 2021-04-23 20:08:54 +07:00
parent a97fdb9548
commit 8353710194
Signed by: blankie
GPG Key ID: CC15FC822C7F61F5
1 changed files with 54 additions and 9 deletions

View File

@ -9,6 +9,8 @@ from io import StringIO, BytesIO
from pyrogram import Client, filters from pyrogram import Client, filters
from .. import config, help_dict, log_errors, slave, apps, session, public_log_errors 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]+)$' 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)) @Client.on_message(~filters.forwarded & ~filters.sticker & ~filters.via_bot & ~filters.edited & filters.me & filters.regex(PYEXEC_REGEX))
@log_errors @log_errors
@ -39,8 +41,15 @@ async def pyexec(client, message):
if ex.msg != "'return' with value in async generator": if ex.msg != "'return' with value in async generator":
raise raise
exx = _gf(obody) exx = _gf(obody)
reply = await message.reply_text('Executing...') reply = await message.reply_text(f'Executing <code>{hash(UniqueExecReturnIdentifier)}</code>...')
async_obj = exx(client, client, message, message, reply, message.reply_to_message, message.reply_to_message, 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 stdout = sys.stdout
stderr = sys.stderr stderr = sys.stderr
wrapped_stdout = StringIO() wrapped_stdout = StringIO()
@ -48,15 +57,19 @@ async def pyexec(client, message):
try: try:
sys.stdout = wrapped_stdout sys.stdout = wrapped_stdout
sys.stderr = wrapped_stderr sys.stderr = wrapped_stderr
if inspect.isasyncgen(async_obj): task = asyncio.create_task(async_obj())
returned = [i async for i in async_obj] exec_tasks[hash(UniqueExecReturnIdentifier)] = task
else: returned = await async_obj
returned = [await async_obj] except asyncio.CancelledError:
if returned == [UniqueExecReturnIdentifier]: sys.stdout = stdout
returned = [] sys.stderr = stderr
exec_tasks.pop(hash(UniqueExecReturnIdentifier), None)
await reply.edit_text('Cancelled')
return
finally: finally:
sys.stdout = stdout sys.stdout = stdout
sys.stderr = stderr sys.stderr = stderr
exec_tasks.pop(hash(UniqueExecReturnIdentifier), None)
wrapped_stderr.seek(0) wrapped_stderr.seek(0)
wrapped_stdout.seek(0) wrapped_stdout.seek(0)
output = '' output = ''
@ -83,4 +96,36 @@ async def pyexec(client, message):
else: else:
await reply.edit_text(output) await reply.edit_text(output)
help_dict['exec'] = ('Exec', '{prefix}exec <i>&lt;python code&gt;</i> - 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'<code>{i}</code>', 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 <i>&lt;python code&gt;</i> - Executes python code
{prefix}listexecs - List exec tasks
Aliases: {prefix}listexec, {prefix}lexec
{prefix}cancelexec <i>&lt;task id&gt;</i> - Cancel exec task
Aliases: {prefix}cexec''')