Skip to content

Commit

Permalink
better crash handling + improved start, stop, reload time
Browse files Browse the repository at this point in the history
  • Loading branch information
@rking32
rking32 committed Mar 11, 2022
1 parent adca6e2 commit fccb79b
Show file tree
Hide file tree
Showing 4 changed files with 67 additions and 60 deletions.
1 change: 1 addition & 0 deletions userge/config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

_LOG = logging.getLogger(__name__)

# try to get this value using eval :)
TEST = secured_str("nice! report @UsergeSpam")

API_ID = environ.get("API_ID")
Expand Down
50 changes: 36 additions & 14 deletions userge/core/types/new/manager.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,15 @@
__all__ = ['Manager']

import asyncio
import logging
from itertools import islice, chain
from typing import Union, List, Dict, Optional

from userge import config
from ..raw import Filter, Command, Plugin
from ... import client as _client, get_collection # pylint: disable=unused-import

_LOG = logging.getLogger(__name__)
_FLT = Union[Filter, Command]


Expand Down Expand Up @@ -235,31 +238,50 @@ def should_wait(self) -> bool:
async def wait(self) -> None:
await self._event.wait()

async def start(self) -> None:
self._event.clear()
async def _do_plugins(self, meth: str) -> None:
loop = asyncio.get_running_loop()
data = iter(self.plugins.values())

for plg in self.plugins.values():
await plg.start()
while True:
chunk = islice(data, config.WORKERS)

self._event.set()
try:
plg = next(chunk)
except StopIteration:
break

async def stop(self) -> None:
self._event.clear()
tasks = []

for plg in self.plugins.values():
await plg.stop()
for plg in chain((plg,), chunk):
tasks.append((plg, loop.create_task(getattr(plg, meth)())))

for plg, task in tasks:
try:
await task
except Exception as i_e:
_LOG.error(f"({meth}) [{plg.cat}/{plg.name}] - {i_e}")

tasks.clear()

_LOG.info(f"on_{meth} tasks completed !")

async def start(self) -> None:
self._event.clear()
await self._do_plugins('start')
self._event.set()

def clear(self) -> None:
for plg in self.plugins.values():
plg.clear()
async def stop(self) -> None:
self._event.clear()
await self._do_plugins('stop')
self._event.set()

async def exit(self) -> None:
await self._do_plugins('exit')
self.plugins.clear()

async def exit(self) -> None:
def clear(self) -> None:
for plg in self.plugins.values():
await plg.exit()
plg.clear()

self.plugins.clear()

Expand Down
18 changes: 1 addition & 17 deletions userge/plugins/builtin/help/__main__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,7 @@ async def helpme(message: Message) -> None: # pylint: disable=missing-function-


if userge.has_bot:

def check_owner(func):
async def wrapper(_, c_q: CallbackQuery):
if c_q.from_user and c_q.from_user.id in config.OWNER_ID:
Expand All @@ -132,7 +133,6 @@ async def wrapper(_, c_q: CallbackQuery):

return wrapper


@userge.bot.on_message(
filters.private & filters.user(list(config.OWNER_ID)) & filters.command("start"), group=-1
)
Expand All @@ -157,7 +157,6 @@ async def pm_help_handler(_, msg: Message):

await msg.reply(out_str, parse_mode='html', disable_web_page_preview=True)


@userge.bot.on_callback_query(filters=filters.regex(pattern=r"\((.+)\)(next|prev)\((\d+)\)"))
@check_owner
async def callback_next_prev(callback_query: CallbackQuery):
Expand All @@ -182,7 +181,6 @@ async def callback_next_prev(callback_query: CallbackQuery):
await callback_query.edit_message_reply_markup(
reply_markup=InlineKeyboardMarkup(buttons))


@userge.bot.on_callback_query(filters=filters.regex(pattern=r"back\((.+)\)"))
@check_owner
async def callback_back(callback_query: CallbackQuery):
Expand All @@ -204,7 +202,6 @@ async def callback_back(callback_query: CallbackQuery):
await callback_query.edit_message_text(
text, reply_markup=InlineKeyboardMarkup(buttons))


@userge.bot.on_callback_query(filters=filters.regex(pattern=r"enter\((.+)\)"))
@check_owner
async def callback_enter(callback_query: CallbackQuery):
Expand All @@ -221,7 +218,6 @@ async def callback_enter(callback_query: CallbackQuery):
await callback_query.edit_message_text(
text, reply_markup=InlineKeyboardMarkup(buttons))


@userge.bot.on_callback_query(
filters=filters.regex(pattern=r"((?:un)?load)\((.+)\)"))
@check_owner
Expand All @@ -247,14 +243,12 @@ async def callback_manage(callback_query: CallbackQuery):
await callback_query.edit_message_text(
text, reply_markup=InlineKeyboardMarkup(buttons))


@userge.bot.on_callback_query(filters=filters.regex(pattern=r"^mm$"))
@check_owner
async def callback_mm(callback_query: CallbackQuery):
await callback_query.edit_message_text(
"🖥 **Userge Main Menu** 🖥", reply_markup=InlineKeyboardMarkup(main_menu_buttons()))


@userge.bot.on_callback_query(filters=filters.regex(pattern=r"^chgclnt$"))
@check_owner
async def callback_chgclnt(callback_query: CallbackQuery):
Expand All @@ -272,7 +266,6 @@ async def callback_chgclnt(callback_query: CallbackQuery):
await callback_query.edit_message_reply_markup(
reply_markup=InlineKeyboardMarkup(main_menu_buttons()))


@userge.bot.on_callback_query(filters=filters.regex(pattern=r"refresh\((.+)\)"))
@check_owner
async def callback_exit(callback_query: CallbackQuery):
Expand All @@ -286,7 +279,6 @@ async def callback_exit(callback_query: CallbackQuery):
await callback_query.edit_message_text(
text, reply_markup=InlineKeyboardMarkup(buttons))


@userge.bot.on_callback_query(filters=filters.regex(pattern=r"prvtmsg\((.+)\)"))
async def prvt_msg(_, c_q: CallbackQuery):
msg_id = str(c_q.matches[0].group(1))
Expand All @@ -303,13 +295,11 @@ async def prvt_msg(_, c_q: CallbackQuery):
await c_q.answer(
f"Only {flname} can see this Private Msg... 😔", show_alert=True)


def is_filter(name: str) -> bool:
split_ = name.split('.')

return bool(split_[0] and len(split_) == 2)


def parse_buttons(page_num: int,
cur_pos: str,
func: Callable[[str], str],
Expand Down Expand Up @@ -339,13 +329,11 @@ def parse_buttons(page_num: int,

return pairs


def main_menu_buttons():
return parse_buttons(0, "mm",
lambda x: f"{_CATEGORY.get(x, '📁')} {x}",
userge.manager.get_all_plugins())


def default_buttons(cur_pos: str):
tmp_btns = []

Expand All @@ -366,7 +354,6 @@ def default_buttons(cur_pos: str):

return [tmp_btns]


def category_data(cur_pos: str):
pos_list = cur_pos.split('|')
plugins = userge.manager.get_all_plugins()[pos_list[1]]
Expand All @@ -380,7 +367,6 @@ def category_data(cur_pos: str):

return text, buttons


def plugin_data(cur_pos: str, p_num: int = 0):

pos_list = cur_pos.split('|')
Expand Down Expand Up @@ -412,7 +398,6 @@ def plugin_data(cur_pos: str, p_num: int = 0):

return text, buttons


def filter_data(cur_pos: str):
pos_list = cur_pos.split('|')
plg = userge.manager.plugins[pos_list[2]]
Expand Down Expand Up @@ -449,7 +434,6 @@ def filter_data(cur_pos: str):

return text, buttons


@userge.bot.on_inline_query(group=1)
async def inline_answer(_, inline_query: InlineQuery):
results = [
Expand Down
58 changes: 29 additions & 29 deletions userge/sys_tools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,35 @@
from typing import Dict, Optional


def secured_env(key: str, default: Optional[str] = None) -> Optional[str]:
""" get secured env """
if not key:
raise ValueError

try:
value = environ.pop(key)
except KeyError:
value = default

ret: Optional[str] = None

if value:
ret = secured_str(value)

return ret


def secured_str(value: str) -> str:
""" get secured string """
if not value:
raise ValueError

ret = _SafeStr(_ST)
ret._ = value

return ret


class SafeDict(Dict[str, str]):
""" modded dict """
def __missing__(self, key: str) -> str:
Expand Down Expand Up @@ -63,33 +92,4 @@ def __str__(self):
return self


def secured_env(key: str, default: Optional[str] = None) -> Optional[str]:
""" get secured env """
if not key:
raise ValueError

try:
value = environ.pop(key)
except KeyError:
value = default

ret: Optional[str] = None

if value:
ret = secured_str(value)

return ret


def secured_str(value: str) -> str:
""" get secured string """
if not value:
raise ValueError

ret = _SafeStr(_ST)
ret._ = value

return ret


_ST = "[SECURED!]"

0 comments on commit fccb79b

Please sign in to comment.