Initial commit

src/cmd.py

@@ -0,0 +1,104 @@
+import argparse
+import asyncio
+import random
+import sys
+from asyncio import Task
+
+from loguru import logger
+from prompt_toolkit import PromptSession, print_formatted_text, ANSI
+from prompt_toolkit.patch_stdout import patch_stdout
+
+from src.adb import Device
+from src.api import get_token
+from src.config import Config
+from src.rip import rip_song, rip_album
+from src.types import GlobalAuthParams
+from src.url import AppleMusicURL, URLType
+
+
+class NewInteractiveShell:
+    loop: asyncio.AbstractEventLoop
+    config: Config
+    tasks: list[Task] = []
+    devices: list[Device] = []
+    storefront_device_mapping: dict[str, list[Device]] = {}
+    anonymous_access_token: str
+    parser: argparse.ArgumentParser
+
+    def __init__(self, loop: asyncio.AbstractEventLoop):
+        self.loop = loop
+        self.config = Config.load_from_config()
+        self.anonymous_access_token = loop.run_until_complete(get_token())
+
+        self.parser = argparse.ArgumentParser(exit_on_error=False)
+        subparser = self.parser.add_subparsers()
+        download_parser = subparser.add_parser("download")
+        download_parser.add_argument("url", type=str)
+        download_parser.add_argument("-c", "--codec",
+                                     choices=["alac", "ec3", "aac", "aac-binaural", "aac-downmix"], default="alac")
+        download_parser.add_argument("-f", "--force", type=bool, default=False)
+        subparser.add_parser("exit")
+
+        logger.remove()
+        logger.add(lambda msg: print_formatted_text(ANSI(msg), end=""), colorize=True, level="INFO")
+
+        for device_info in self.config.devices:
+            device = Device(frida_path=device_info.fridaPath)
+            device.connect(device_info.host, device_info.port)
+            logger.info(f"Device {device_info.host}:{device_info.port} has connected")
+            self.devices.append(device)
+            auth_params = device.get_auth_params()
+            if not self.storefront_device_mapping.get(auth_params.storefront.lower()):
+                self.storefront_device_mapping.update({auth_params.storefront.lower(): []})
+            self.storefront_device_mapping[auth_params.storefront.lower()].append(device)
+            device.start_inject_frida(device_info.agentPort)
+
+    async def command_parser(self, cmd: str):
+        if not cmd.strip():
+            return
+        cmds = cmd.split(" ")
+        try:
+            args = self.parser.parse_args(cmds)
+        except argparse.ArgumentError:
+            logger.warning(f"Unknown command: {cmd}")
+            return
+        match cmds[0]:
+            case "download":
+                await self.do_download(args.url, args.codec, args.force)
+            case "exit":
+                self.loop.stop()
+                sys.exit()
+
+    async def do_download(self, raw_url: str, codec: str, force_download: bool):
+        url = AppleMusicURL.parse_url(raw_url)
+        devices = self.storefront_device_mapping.get(url.storefront)
+        if not devices:
+            logger.error(f"No device is available to decrypt the specified region: {url.storefront}")
+        available_devices = [device for device in devices if not device.decryptLock.locked()]
+        if not available_devices:
+            available_device: Device = random.choice(devices)
+        else:
+            available_device: Device = random.choice(available_devices)
+        global_auth_param = GlobalAuthParams.from_auth_params_and_token(available_device.get_auth_params(), self.anonymous_access_token)
+        match url.type:
+            case URLType.Song:
+                self.loop.create_task(rip_song(url, global_auth_param, codec, self.config, available_device, force_download))
+            case URLType.Album:
+                self.loop.create_task(rip_album(url, global_auth_param, codec, self.config, available_device))
+
+    async def handle_command(self):
+        session = PromptSession("> ")
+
+        while True:
+            try:
+                command = await session.prompt_async()
+                await self.command_parser(command)
+            except (EOFError, KeyboardInterrupt):
+                return
+
+    async def start(self):
+        with patch_stdout():
+            try:
+                await self.handle_command()
+            finally:
+                logger.info("Existing shell")