diff --git a/modules/decorators/callback.py b/modules/decorators/callback.py index df71c14..19d04e1 100644 --- a/modules/decorators/callback.py +++ b/modules/decorators/callback.py @@ -1,10 +1,17 @@ +from rich import print + from modules.menu.menu import menu +from ..exceptions.pretty_exception import PrettyException def callback(f): def wrapper(*args, **kwargs): menu.pause() - f(*args, **kwargs) + try: + f(*args, **kwargs) + except Exception as e: + print(PrettyException(e).pretty_exception) + input() menu.resume() return wrapper @@ -14,7 +21,11 @@ def async_callback(f): def wrapper(*args, **kwargs): import asyncio menu.pause() - asyncio.run(f(*args, **kwargs)) + try: + asyncio.run(f(*args, **kwargs)) + except Exception as e: + print(PrettyException(e).pretty_exception) + input() menu.resume() return wrapper diff --git a/modules/exceptions/__init__.py b/modules/exceptions/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/modules/exceptions/pretty_exception.py b/modules/exceptions/pretty_exception.py new file mode 100644 index 0000000..d09a6a9 --- /dev/null +++ b/modules/exceptions/pretty_exception.py @@ -0,0 +1,47 @@ +import os +import traceback +import contextlib +import re + + +class PrettyException: + def __init__(self, e: Exception): + self.pretty_exception = f'❌ Error! \n' \ + f'🐊 {e.__traceback__.tb_frame.f_code.co_filename.replace(os.getcwd(), "")}' \ + f':{e.__traceback__.tb_frame.f_lineno} \n' \ + f'😍 {e.__class__.__name__} \n' \ + f'👉 {"".join(traceback.format_exception_only(e)).strip()} \n\n' \ + f'⬇️ Trace: \n' \ + f'{self.get_full_stack()}' + + @staticmethod + def get_full_stack(): + full_stack = traceback.format_exc().replace( + "Traceback (most recent call last):\n", "" + ) + + line_regex = r' File "(.*?)", line ([0-9]+), in (.+)' + + def format_line(line: str) -> str: + filename_, lineno_, name_ = re.search(line_regex, line).groups() + with contextlib.suppress(Exception): + filename_ = os.path.basename(filename_) + + return ( + f"🤯 {filename_}:{lineno_} (in" + f" {name_} call)" + ) + + full_stack = "\n".join( + [ + format_line(line) + if re.search(line_regex, line) + else f"{line}" + for line in full_stack.splitlines() + ] + ) + + return full_stack + + def __str__(self): + return self.pretty_exception diff --git a/modules/menu/callbacks/create_session/callbacks.py b/modules/menu/callbacks/create_session/callbacks.py index ce9a75f..cb20db6 100644 --- a/modules/menu/callbacks/create_session/callbacks.py +++ b/modules/menu/callbacks/create_session/callbacks.py @@ -21,6 +21,7 @@ async def create_session_callback(): username=client.me.username ).json() print(f'[green]Created[/] session {session_name}...') + await client.stop() input() except OperationalError: diff --git a/modules/menu/callbacks/get_code/__init__.py b/modules/menu/callbacks/get_code/__init__.py index a2ae2a9..0d4875e 100644 --- a/modules/menu/callbacks/get_code/__init__.py +++ b/modules/menu/callbacks/get_code/__init__.py @@ -1 +1 @@ -from .get_code import get_code_callback +from .callbacks import get_code_callback diff --git a/modules/menu/callbacks/get_code/get_code.py b/modules/menu/callbacks/get_code/callbacks.py similarity index 100% rename from modules/menu/callbacks/get_code/get_code.py rename to modules/menu/callbacks/get_code/callbacks.py diff --git a/modules/menu/callbacks/update_info/__init__.py b/modules/menu/callbacks/update_info/__init__.py new file mode 100644 index 0000000..0c2d849 --- /dev/null +++ b/modules/menu/callbacks/update_info/__init__.py @@ -0,0 +1 @@ +from .callbacks import update_info_callback diff --git a/modules/menu/callbacks/update_info/callbacks.py b/modules/menu/callbacks/update_info/callbacks.py new file mode 100644 index 0000000..67fb56b --- /dev/null +++ b/modules/menu/callbacks/update_info/callbacks.py @@ -0,0 +1,22 @@ +from modules.decorators.callback import async_callback +from modules.client import GeneratedClient +from modules.config import sessions +from modules.config.models import SessionConfig +from rich import print + + +@async_callback +async def update_info_callback(session_name: str): + print('Getting info about profile...') + await (client := GeneratedClient(name=session_name)).start() + + sessions[session_name] = SessionConfig( + id=client.me.id, + phone=f'+{client.me.phone_number}', + profile_name=client.me.first_name + (f' {client.me.last_name}' if client.me.last_name else ''), + username=client.me.username + ).json() + print(f'[green]Updated[/] info for {session_name}. Go back in menu to update it in interface...') + + await client.stop() + input() diff --git a/modules/menu/dynamic/common_generators/__init__.py b/modules/menu/dynamic/common_generators/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/modules/menu/dynamic/get_code/generator.py b/modules/menu/dynamic/common_generators/valid_sessions_generator.py similarity index 67% rename from modules/menu/dynamic/get_code/generator.py rename to modules/menu/dynamic/common_generators/valid_sessions_generator.py index af8e14f..b0daec7 100644 --- a/modules/menu/dynamic/get_code/generator.py +++ b/modules/menu/dynamic/common_generators/valid_sessions_generator.py @@ -1,21 +1,22 @@ import os from cursesmenu import CursesMenu from cursesmenu.items import FunctionItem +from typing import Callable, Any from modules.config import sessions from modules.config.models import SessionConfig -from ...callbacks.get_code import get_code_callback -def generate_get_code_menu() -> CursesMenu: +def generate_sessions_list(menu_name: str, callback: Callable[..., Any]) -> CursesMenu: submenu = CursesMenu( - title='Sessions' + title=menu_name, + subtitle='name - phone - id - profile name - username' ) for session_name in sessions.keys(): if os.path.isfile(f'{session_name}.session'): submenu.items.append(FunctionItem( f'{session_name} - {SessionConfig(**sessions[session_name])}', - function=get_code_callback, + function=callback, args=[session_name] )) diff --git a/modules/menu/dynamic/delete/__init__.py b/modules/menu/dynamic/delete/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/modules/menu/dynamic/get_code/submenu.py b/modules/menu/dynamic/get_code/submenu.py index 1a0a347..6e0cd4f 100644 --- a/modules/menu/dynamic/get_code/submenu.py +++ b/modules/menu/dynamic/get_code/submenu.py @@ -1,7 +1,8 @@ from ...menu import menu from cursesmenu import CursesMenu from modules.custom_items.DynamicSubmenu import DynamicSubmenuItem -from .generator import generate_get_code_menu +from ..common_generators.valid_sessions_generator import generate_sessions_list +from ...callbacks.get_code import get_code_callback get_code_submenu = CursesMenu( @@ -10,6 +11,7 @@ get_code_submenu = CursesMenu( get_code_submenu_item = DynamicSubmenuItem( text='Get confirmation code', - generator=generate_get_code_menu, + generator=generate_sessions_list, + args=['Get code for', get_code_callback], menu=menu ) diff --git a/modules/menu/dynamic/update/__init__.py b/modules/menu/dynamic/update/__init__.py new file mode 100644 index 0000000..3c1382e --- /dev/null +++ b/modules/menu/dynamic/update/__init__.py @@ -0,0 +1 @@ +from .submenu import get_submenu_item diff --git a/modules/menu/dynamic/update/submenu.py b/modules/menu/dynamic/update/submenu.py new file mode 100644 index 0000000..32a1045 --- /dev/null +++ b/modules/menu/dynamic/update/submenu.py @@ -0,0 +1,13 @@ +from cursesmenu import CursesMenu +from modules.custom_items.DynamicSubmenu import DynamicSubmenuItem +from ..common_generators.valid_sessions_generator import generate_sessions_list +from ...callbacks.update_info import update_info_callback + + +def get_submenu_item(parent_menu: CursesMenu): + return DynamicSubmenuItem( + text='Update account info', + generator=generate_sessions_list, + args=['Update info for', update_info_callback], + menu=parent_menu + ) diff --git a/modules/menu/items/main.py b/modules/menu/items/main.py index 197d9e7..a13885b 100644 --- a/modules/menu/items/main.py +++ b/modules/menu/items/main.py @@ -1,10 +1,12 @@ from .config import config_submenu_item from .create_session import create_new_session_item from ..dynamic.get_code import get_code_submenu_item +from .manage import manage_submenu_item items_list = [ create_new_session_item, get_code_submenu_item, + manage_submenu_item, config_submenu_item ] diff --git a/modules/menu/items/manage/__init__.py b/modules/menu/items/manage/__init__.py new file mode 100644 index 0000000..eed7343 --- /dev/null +++ b/modules/menu/items/manage/__init__.py @@ -0,0 +1 @@ +from .submenu import manage_submenu_item diff --git a/modules/menu/items/manage/fields.py b/modules/menu/items/manage/fields.py new file mode 100644 index 0000000..b0941e8 --- /dev/null +++ b/modules/menu/items/manage/fields.py @@ -0,0 +1,9 @@ +from cursesmenu import CursesMenu +from ...dynamic import update + + +def get_items_list(menu: CursesMenu): + _ = [ + update.get_submenu_item + ] + return [x(menu) for x in _] diff --git a/modules/menu/items/manage/submenu.py b/modules/menu/items/manage/submenu.py new file mode 100644 index 0000000..e4bb2d2 --- /dev/null +++ b/modules/menu/items/manage/submenu.py @@ -0,0 +1,17 @@ +from ...menu import menu +from cursesmenu import CursesMenu +from cursesmenu.items import SubmenuItem +from .fields import get_items_list + + +manage_submenu = CursesMenu( + title='Manage' +) +for item in get_items_list(manage_submenu): + manage_submenu.items.append(item) + +manage_submenu_item = SubmenuItem( + text='Manage/edit/delete sessions', + submenu=manage_submenu, + menu=menu +)