IPAToolClient/main.py

import asyncio
import getpass
import traceback
from pathlib import Path
from shutil import move

from aiohttp import ClientTimeout
from httpwrapper import AsyncClientConfig, BaseAsyncClient
from jsoner import json_read_async, json_write_async

from models import UdidInfo

CACHE_DIR = Path("cache")
CACHE_DIR.mkdir(exist_ok=True)

APPS_DIR = Path("apps")
APPS_DIR.mkdir(exist_ok=True)


def move_ipa(from_path: str, to_path: str):
    while True:
        try:
            move(from_path, to_path)
            return
        except Exception:
            traceback.print_exc()


class IPAToolClient(BaseAsyncClient):
    def __init__(self, host: str, port: int):
        super().__init__(
            f"http://{host}:{port}",
            config=AsyncClientConfig(timeout=ClientTimeout(300)),
        )

    async def auth(self, user: str, pswd: str) -> bool:
        r = await self._post(
            "/auth/login",
            params={
                "user": user,
                "pswd": pswd,
                "keychain": pswd,
            },
        )
        if r.status != 200:
            return False
        json_data = await r.json()
        if isinstance(json_data, bool):
            return json_data
        return False

    async def auth_check(self, user: str) -> bool:
        r = await self._post(
            "/auth/check",
            params={
                "user": user,
            },
        )
        if r.status != 200:
            return False
        json_data = await r.json()
        if isinstance(json_data, bool):
            return json_data
        return False

    async def app_info(self, query: str, limit: int, keychain: str) -> dict:
        r = await self._post(
            "/app/search",
            params={
                "query": query,
                "limit": limit,
                "keychain": keychain,
            },
        )
        if r.status != 200:
            return {}
        json_data = await r.json()
        if isinstance(json_data, dict):
            return json_data
        return {}

    async def app_upload(self, output_path: str, save_path: str):
        r = await self._get(
            "/app/upload",
            params={
                "pswd": "81928192",
                "output_path": output_path,
            },
        )
        if r.status != 200:
            return False
        try:
            content = await r.read()
            with open(save_path, "wb") as f:
                f.write(content)
            return True
        except Exception:
            traceback.print_exc()
            return False

    async def app_download(
        self,
        bundle_id: str,
        output_path: str,
        keychain: str,
    ) -> bool:
        r = await self._post(
            "/app/download",
            params={
                "keychain": keychain,
                "bundle_id": bundle_id,
                "output_path": output_path,
            },
        )
        if r.status != 200:
            return False
        json_data = await r.json()
        if isinstance(json_data, dict):
            return json_data.get("success", False)
        return False

    async def app_delete(self, output_path: str) -> bool:
        r = await self._delete(
            "/app/delete",
            params={
                "pswd": "81928192",
                "output_path": output_path,
            },
        )
        if r.status != 200:
            return False
        json_data = await r.json()
        if isinstance(json_data, bool):
            return json_data
        return False


async def __main(client: IPAToolClient, udid: str, info: UdidInfo):
    output_path = f"{info.app_name}.ipa"
    await client.app_delete(output_path)

    for _ in range(30):
        if await client.app_download(info.bundle_id, output_path, info.pswd):
            break
        print(udid, "download failed, retrying...")
        await asyncio.sleep(3)
    print(udid, "downloaded")

    save_path = APPS_DIR / f"{info.app_name}-{udid}.ipa"
    for _ in range(10):
        if await client.app_upload(output_path, str(save_path)):
            return True
    return False


async def _main(client: IPAToolClient, udid: str, info: UdidInfo) -> bool:
    if not await client.auth(info.user, info.pswd):
        print(udid, "no auth response")
        return False
    if not (response := await client.app_info(info.app_name, 1, info.pswd)):
        print(udid, "no app info response")
        return False
    if not (apps := response.get("apps", [])):
        print(udid, "no apps in response")
        return False
    for app in apps:
        if app.get("bundleID", "") != info.bundle_id:
            continue
        if not (version := app.get("version", "")):
            print(udid, "no version in response")
            continue
        old_version = ""
        cache_file = CACHE_DIR / f"{udid}.{info.app_name}.json"
        if json_data := await json_read_async(cache_file):
            old_version = json_data.get("version", "")
        if old_version == version:
            print(udid, "already have version", version)
            return False
        print(udid, "new version", version)
        # await json_write_async(cache_file, {"version": version})
        if await __main(client, udid, info):
            await json_write_async(cache_file, {"version": version})

            print(udid, "uploaded")

            ipa_file = f"{info.app_name}-{udid}.ipa"
            save_path = APPS_DIR / ipa_file
            to_path = f"/Users/{getpass.getuser()}/{ipa_file}"
            move_ipa(str(save_path), to_path)
            print(udid, "moved")
            return True
        print(udid, "not uploaded")
        return False
    return False


async def main(json_data: dict):
    tasks = []
    for udid, value in json_data.items():
        info = UdidInfo(**value)
        client = IPAToolClient(info.host, info.port)
        task = _main(client, udid, info)
        tasks.append(task)
    await asyncio.gather(*tasks, return_exceptions=True)


async def start():
    while True:
        json_data = await json_read_async("info.json")
        await main(json_data)
        await asyncio.sleep(120)


if __name__ == "__main__":
    asyncio.run(start())