Merge pull request #49 from itouakirai/v2

type and feat

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: