From a882f370bbfa74317c051ea06fe56596b4724297 Mon Sep 17 00:00:00 2001 From: rking32 Date: Sat, 17 Oct 2020 19:38:11 +0530 Subject: [PATCH] some fixes + security patch (sensitive vars now secured) --- runtime.txt | 2 +- userge/config.py | 1 - userge/core/client.py | 33 +++++++++---------- .../methods/messages/edit_message_text.py | 3 +- userge/core/methods/messages/send_as_file.py | 3 +- userge/core/methods/messages/send_message.py | 3 +- userge/plugins/tools/executor.py | 11 ++----- userge/utils/__init__.py | 2 +- userge/utils/sys_tools.py | 20 +++++++++++ 9 files changed, 46 insertions(+), 32 deletions(-) diff --git a/runtime.txt b/runtime.txt index 1124509c8..6f98a22cc 100644 --- a/runtime.txt +++ b/runtime.txt @@ -1 +1 @@ -python-3.8.5 \ No newline at end of file +python-3.8.6 \ No newline at end of file diff --git a/userge/config.py b/userge/config.py index 5082b432b..6b689ceb2 100644 --- a/userge/config.py +++ b/userge/config.py @@ -47,7 +47,6 @@ class Config: INSTA_PASS = os.environ.get("INSTA_PASS") UPSTREAM_REPO = os.environ.get("UPSTREAM_REPO") UPSTREAM_REMOTE = os.environ.get("UPSTREAM_REMOTE") - SCREENSHOT_API = os.environ.get("SCREENSHOT_API", None) SPAM_WATCH_API = os.environ.get("SPAM_WATCH_API", None) CURRENCY_API = os.environ.get("CURRENCY_API", None) OCR_SPACE_API_KEY = os.environ.get("OCR_SPACE_API_KEY", None) diff --git a/userge/core/client.py b/userge/core/client.py index 37e71926d..f8570b246 100644 --- a/userge/core/client.py +++ b/userge/core/client.py @@ -77,8 +77,7 @@ async def load_plugin(self, name: str, reload_plugin: bool = False) -> None: if hasattr(plg, '_init'): # pylint: disable=protected-access if asyncio.iscoroutinefunction(plg._init): - _INIT_TASKS.append( - asyncio.get_event_loop().create_task(plg._init())) + _INIT_TASKS.append(self.loop.create_task(plg._init())) _LOG.debug(_LOG_STR, f"Imported {_IMPORTED[-1].__name__} Plugin Successfully") async def _load_plugins(self) -> None: @@ -173,30 +172,28 @@ async def stop(self) -> None: # pylint: disable=arguments-differ def begin(self, coro: Optional[Awaitable[Any]] = None) -> None: """ start userge """ - loop = asyncio.get_event_loop() - loop.add_signal_handler(signal.SIGHUP, _shutdown) - loop.add_signal_handler(signal.SIGTERM, _shutdown) - run = loop.run_until_complete + self.loop.add_signal_handler(signal.SIGHUP, _shutdown) + self.loop.add_signal_handler(signal.SIGTERM, _shutdown) + run = self.loop.run_until_complete + run(self.start()) + running_tasks: List[asyncio.Task] = [] + for task in self._tasks: + running_tasks.append(self.loop.create_task(task())) + logbot.edit_last_msg("Userge has Started Successfully !") + logbot.end() try: - run(self.start()) - running_tasks: List[asyncio.Task] = [] - for task in self._tasks: - running_tasks.append(loop.create_task(task())) if coro: _LOG.info(_LOG_STR, "Running Coroutine") run(coro) else: _LOG.info(_LOG_STR, "Idling Userge") - logbot.edit_last_msg("Userge has Started Successfully !") - logbot.end() idle() + except asyncio.exceptions.CancelledError: + pass + finally: _LOG.info(_LOG_STR, "Exiting Userge") for task in running_tasks: task.cancel() run(self.stop()) - run(loop.shutdown_asyncgens()) - except asyncio.exceptions.CancelledError: - pass - finally: - if not loop.is_running(): - loop.close() + run(self.loop.shutdown_asyncgens()) + self.loop.close() diff --git a/userge/core/methods/messages/edit_message_text.py b/userge/core/methods/messages/edit_message_text.py index 82b3ce415..e7c22a76d 100644 --- a/userge/core/methods/messages/edit_message_text.py +++ b/userge/core/methods/messages/edit_message_text.py @@ -17,6 +17,7 @@ from pyrogram.types import InlineKeyboardMarkup from userge import Config +from userge.utils import secure_text from ...ext import RawClient from ... import types @@ -84,7 +85,7 @@ async def edit_message_text(self, # pylint: disable=arguments-differ """ msg = await super().edit_message_text(chat_id=chat_id, message_id=message_id, - text=text, + text=secure_text(text), parse_mode=parse_mode, disable_web_page_preview=disable_web_page_preview, reply_markup=reply_markup) diff --git a/userge/core/methods/messages/send_as_file.py b/userge/core/methods/messages/send_as_file.py index 993ebffe4..2f09633a0 100644 --- a/userge/core/methods/messages/send_as_file.py +++ b/userge/core/methods/messages/send_as_file.py @@ -17,6 +17,7 @@ import aiofiles from userge import logging +from userge.utils import secure_text from ...ext import RawClient from ... import types @@ -70,7 +71,7 @@ async def send_as_file(self, On success, the sent Message is returned. """ async with aiofiles.open(filename, "w+", encoding="utf8") as out_file: - await out_file.write(text) + await out_file.write(secure_text(text)) _LOG.debug(_LOG_STR, f"Uploading {filename} To Telegram") msg = await self.send_document(chat_id=chat_id, document=filename, diff --git a/userge/core/methods/messages/send_message.py b/userge/core/methods/messages/send_message.py index e304cce2c..33ffc8c81 100644 --- a/userge/core/methods/messages/send_message.py +++ b/userge/core/methods/messages/send_message.py @@ -19,6 +19,7 @@ ReplyKeyboardRemove, ForceReply) from userge import Config +from userge.utils import secure_text from ...ext import RawClient from ... import types @@ -92,7 +93,7 @@ async def send_message(self, # pylint: disable=arguments-differ :obj:`Message`: On success, the sent text message or True is returned. """ msg = await super().send_message(chat_id=chat_id, - text=text, + text=secure_text(text), parse_mode=parse_mode, disable_web_page_preview=disable_web_page_preview, disable_notification=disable_notification, diff --git a/userge/plugins/tools/executor.py b/userge/plugins/tools/executor.py index ecf8c0d27..c3f087481 100644 --- a/userge/plugins/tools/executor.py +++ b/userge/plugins/tools/executor.py @@ -16,8 +16,6 @@ from getpass import getuser from os import geteuid -from pyrogram.errors.exceptions.bad_request_400 import MessageNotModified - from userge import userge, Message, Config from userge.utils import runcmd @@ -71,7 +69,7 @@ async def aexec(code): output = "" if not silent_mode: output += f"**>** ```{cmd}```\n\n" - if evaluation: + if evaluation is not None: output += f"**>>** ```{evaluation}```" if output: await message.edit_or_send_as_file(text=output, @@ -143,11 +141,8 @@ async def term_(message: Message): out_data = f"
{output}{t_obj.read_line}
" await message.try_to_edit(out_data, parse_mode='html') out_data = f"
{output}{t_obj.get_output}
" - try: - await message.edit_or_send_as_file( - out_data, parse_mode='html', filename="term.txt", caption=cmd) - except MessageNotModified: - pass + await message.edit_or_send_as_file( + out_data, parse_mode='html', filename="term.txt", caption=cmd) async def init_func(message: Message): diff --git a/userge/utils/__init__.py b/userge/utils/__init__.py index e91aadd7e..5989f84a7 100644 --- a/userge/utils/__init__.py +++ b/userge/utils/__init__.py @@ -9,7 +9,7 @@ # All rights reserved. from .progress import progress # noqa -from .sys_tools import SafeDict, get_import_path # noqa +from .sys_tools import SafeDict, get_import_path, secure_text # noqa from .tools import (demojify, # noqa get_file_id_and_ref, humanbytes, diff --git a/userge/utils/sys_tools.py b/userge/utils/sys_tools.py index f2716f23e..7f17d9a1a 100644 --- a/userge/utils/sys_tools.py +++ b/userge/utils/sys_tools.py @@ -9,9 +9,20 @@ # All rights reserved. from glob import glob +from os import environ from os.path import isfile, relpath from typing import Dict, List, Union +_SECURE = [ + # critical + 'API_ID', 'API_HASH', 'BOT_TOKEN', 'HU_STRING_SESSION', 'DATABASE_URL', 'HEROKU_API_KEY', + # others + 'INSTA_ID', 'INSTA_PASS', 'SPAM_WATCH_API', 'CURRENCY_API', 'OCR_SPACE_API_KEY', + 'REMOVE_BG_API_KEY', 'G_DRIVE_CLIENT_ID', 'G_DRIVE_CLIENT_SECRET', + # unofficial + 'ARL_TOKEN', 'GCS_API_KEY', 'GCS_IMAGE_E_ID', 'G_PHOTOS_CLIENT_ID', + 'G_PHOTOS_CLIENT_SECRET', 'CH_LYDIA_API'] + class SafeDict(Dict[str, str]): """ modded dict """ @@ -31,3 +42,12 @@ def get_import_path(root: str, path: str) -> Union[str, List[str]]: if not f.endswith("__init__.py") ] ) + + +def secure_text(text: str) -> str: + """ secure given text """ + for var in _SECURE: + tvar = environ.get(var, None) + if tvar and tvar in text: + text = text.replace(tvar, "**SECURED!**") + return text