diff --git a/src/cmd.py b/src/cmd.py index 64d90f4..c944dc1 100644 --- a/src/cmd.py +++ b/src/cmd.py @@ -5,11 +5,12 @@ import sys from creart import it from prompt_toolkit import PromptSession from prompt_toolkit.patch_stdout import patch_stdout +from prompt_toolkit.completion import WordCompleter from src.api import WebAPI from src.config import Config from src.flags import Flags -from src.grpc.manager import WrapperManager +from src.grpc.manager import WrapperManager, WrapperManagerException from src.logger import GlobalLogger from src.measurer import SpeedMeasurer from src.rip import on_decrypt_success, on_decrypt_failed, rip_song, rip_album, rip_artist, rip_playlist @@ -49,6 +50,8 @@ class InteractiveShell: download_parser.add_argument("--include-participate-songs", default=False, dest="include", action="store_true") subparser.add_parser("status") + subparser.add_parser("login") + subparser.add_parser("logout") subparser.add_parser("exit") async def show_status(self): @@ -96,18 +99,57 @@ class InteractiveShell: return def bottom_toolbar(self): - return f"Download Speed: {it(SpeedMeasurer).download_speed()}, Decrypt Speed: {it(SpeedMeasurer).decrypt_speed()}, Tasks: {get_tasks_num()}" + return f"Download Speed: {it(SpeedMeasurer).download_speed()}, Decrypt Speed: {it(SpeedMeasurer).decrypt_speed()}, Tasks: {get_tasks_num()-2}" + + def completer(self): + mycompleter = ['dl', 'status', 'login', 'logout', 'exit'] + return WordCompleter(mycompleter) async def handle_command(self): - session = PromptSession("> ", bottom_toolbar=self.bottom_toolbar, refresh_interval=1) + session = PromptSession("> ", bottom_toolbar=self.bottom_toolbar, completer=self.completer(), refresh_interval=1) while True: try: command = await session.prompt_async() - await self.command_parser(command) + if command.lower() == 'login': + await self.login_flow() + if command.lower() == 'logout': + await self.logout_flow() + elif command.strip() == '': + continue + else: await self.command_parser(command) except (EOFError, KeyboardInterrupt): return + async def on_2fa(self, username: str, password: str): + session = PromptSession() + two_step_code = await session.prompt_async("2FA code: ") + return two_step_code + + async def login_flow(self): + await it(WrapperManager).init(it(Config).instance.url, it(Config).instance.secure) + session = PromptSession() + username = await session.prompt_async("Username: ") + password = await session.prompt_async("Password: ", is_password=True) + try: + await it(WrapperManager).login(username, password, self.on_2fa) + except WrapperManagerException as e: + it(GlobalLogger).logger.error("Login Failed!") + return + it(GlobalLogger).logger.info("Login Success!") + + async def logout_flow(self): + await it(WrapperManager).init(it(Config).instance.url, it(Config).instance.secure) + session = PromptSession() + username = await session.prompt_async("Username: ") + try: + await it(WrapperManager).logout(username) + except WrapperManagerException as e: + it(GlobalLogger).logger.error("Logout Failed!") + return + it(GlobalLogger).logger.info("Logout Success!") + + async def start(self): with patch_stdout(): try: