IPAToolClient/main.py

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

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()
            sleep(10)


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())