feat?: 试图加入自定义协议入库

main.py

@@ -1,64 +1,205 @@
 import sys
 import os
+import winreg
 import asyncio
-import re
-
-from colorama import Fore, Back, Style
-from colorama import init as cinit
+from typing import Any
+from aiohttp import ClientSession, ConnectionTimeoutError
+from common.checkcn import checkcn
+from common.config import config
+from common.dkey_merge import depotkey_merge
+from common.migration import migrate
+from common.unlock import stool_add, greenluma_add
+from common.get_manifest_info import get_manifest
+from common.check import check_github_api_rate_limit
 from common.log import log
+from common.get_steam_path import steam_path
 from common.stack_error import stack_error
-from common.init_text import init
-from common.main_func import main
 
-# 摘自https://github.com/MeoProject/lx-music-api-server
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
 lock = asyncio.Lock()
+isGreenLuma = any((steam_path / dll).exists()
+                  for dll in ['GreenLuma_2024_x86.dll', 'GreenLuma_2024_x64.dll', 'User32.dll'])
+isSteamTools = (steam_path / 'config' / 'stUI').is_dir()
 
-init()
-cinit()
 
-repos = [
-    'ikun0014/ManifestHub',
-    'Auiowu/ManifestAutoUpdate',
-    'tymolu233/ManifestAutoUpdate',
-]
+def init():
+    banner_lines = [
+        f"  _____   __   _   _____   _   _    _____  __    __ ",
+        f" /  _  \\ |  \\ | | | ____| | | / /  | ____| \\ \\  / /",
+        f" | | | | |   \\| | | |__   | |/ /   | |__    \\ \\/ /",
+        f" | | | | | |\\   | |  __|  | |\\ \\   |  __|    \\  / ",
+        f" | |_| | | | \\  | | |___  | | \\ \\  | |___    / /",
+        f" \\_____/ |_|  \\_| |_____| |_|  \\_\\ |_____|  /_/",
+    ]
+    for line in banner_lines:
+        log.info(line)
 
+    log.info(f'作者: ikun0014')
+    log.warning(f'本项目采用GNU General Public License v3开源许可证，请勿用于商业用途')
+    log.info(f'版本: 1.3.3')
+    log.info(f'项目Github仓库: https://github.com/ikunshare/Onekey \n Gitee: https://gitee.com/ikun0014/Onekey')
+    log.info(f'官网: ikunshare.com')
+    log.warning(
+        f'本项目完全开源免费, 如果你在淘宝, QQ群内通过购买方式获得, 赶紧回去骂商家死全家\n   交流群组:\n    https://t.me/ikunshare_qun')
 
-def prompt_app_id():
-    while True:
-        app_id = input(f"{Fore.CYAN}{Back.BLACK}{
-                       Style.BRIGHT}请输入游戏AppID: {Style.RESET_ALL}").strip()
-        if re.match(r'^\d+$', app_id):
-            return app_id
-        else:
-            print(f"{Fore.RED}无效的AppID, 请输入数字!{Style.RESET_ALL}")
 
+program_path = sys.argv[0]
+program_name = os.path.basename(program_path)
 
-async def main_loop():
-    while True:
+
+def check_registry_key_exists() -> bool:
+    try:
+        key = winreg.OpenKey(winreg.HKEY_CLASSES_ROOT, "")
+        try:
+            winreg.OpenKey(key, "onekey")
+            return True
+        except FileNotFoundError:
+            return False
+        finally:
+            winreg.CloseKey(key)
+    except FileNotFoundError:
+        return False
+
+
+def create_registry_key_with_values(key_path, default_value=None):
+    try:
+        key = winreg.CreateKey(winreg.HKEY_CLASSES_ROOT, key_path)
+        if default_value is not None:
+            winreg.SetValue(key, '', winreg.REG_SZ, default_value)
+    except Exception as e:
+        log.error(f"协议注册失败: {e}")
+    finally:
         try:
-            app_id = prompt_app_id()
-            await main(app_id, repos)
-        except KeyboardInterrupt:
-            log.info("程序已退出")
-        except EOFError:
-            break
+            winreg.CloseKey(key)
+        except:
+            pass
+
+
+sub_key_paths = [
+    "onekey\\shell",
+    "onekey\\shell\\open",
+    "onekey\\shell\\open\\command"
+]
+command_default_value = f'"{program_name}" app "%1"'
+
+if not check_registry_key_exists():
+    create_registry_key_with_values("onekey", "URL:Onekey Protocol")
+    key = winreg.CreateKey(winreg.HKEY_CLASSES_ROOT, "onekey")
+    winreg.SetValueEx(key, "URL Protocol", 0, winreg.REG_SZ, "")
+    winreg.CloseKey(key)
 
+for i, sub_key_path in enumerate(sub_key_paths):
+    if i == len(sub_key_paths) - 1:
+        create_registry_key_with_values(sub_key_path, command_default_value)
+    else:
+        create_registry_key_with_values(sub_key_path)
 
-async def run():
+async def fetch_branch_info(session, url, headers) -> str | None:
     try:
-        log.info('App ID可以在SteamDB或Steam商店链接页面查看')
-        await main_loop()
+        async with session.get(url, headers=headers, ssl=False) as response:
+            return await response.json()
     except KeyboardInterrupt:
         log.info("程序已退出")
     except Exception as e:
-        log.error(f'发生错误: {stack_error(e)}, 将在5秒后退出')
-        await asyncio.sleep(5)
+        log.error(f'获取信息失败: {stack_error(e)}')
+        return None
+    except ConnectionTimeoutError as e:
+        log.error(f'获取信息时超时: {stack_error(e)}')
+        return None
+
+
+async def get_latest_repo_info(session, repos, app_id, headers) -> Any | None:
+    latest_date = None
+    selected_repo = None
+
+    for repo in repos:
+        url = f'https://api.github.com/repos/{repo}/branches/{app_id}'
+        r_json = await fetch_branch_info(session, url, headers)
+        if r_json and 'commit' in r_json:
+            date = r_json['commit']['commit']['author']['date']
+            if latest_date is None or date > latest_date:
+                latest_date = date
+                selected_repo = repo
+
+    return selected_repo, latest_date
+
+
+async def main(app_id: str) -> bool:
+    repos = [
+        'ikun0014/ManifestHub',
+        'Auiowu/ManifestAutoUpdate',
+        'tymolu233/ManifestAutoUpdate',
+    ]
+
+    app_id_list = list(filter(str.isdecimal, app_id.strip().split('-')))
+
+    if not app_id_list:
+        log.error(f'App ID无效')
+        return False
+
+    app_id = app_id_list[0]
+
+    async with ClientSession() as session:
+        github_token = config.get("Github_Personal_Token", "")
+        headers = {'Authorization': f'Bearer {
+            github_token}'} if github_token else None
+
+        await checkcn(session)
+        await check_github_api_rate_limit(headers, session)
+
+        selected_repo, latest_date = await get_latest_repo_info(session, repos, app_id, headers)
+
+        if selected_repo:
+            log.info(f'选择清单仓库: {selected_repo}')
+            url = f'https://api.github.com/repos/{
+                selected_repo}/branches/{app_id}'
+            r_json = await fetch_branch_info(session, url, headers)
+
+            if r_json and 'commit' in r_json:
+                sha = r_json['commit']['sha']
+                url = r_json['commit']['commit']['tree']['url']
+                r2_json = await fetch_branch_info(session, url, headers)
+
+                if r2_json and 'tree' in r2_json:
+                    collected_depots = []
+                    for item in r2_json['tree']:
+                        result = await get_manifest(sha, item['path'], steam_path, selected_repo, session)
+                        collected_depots.extend(result)
+
+                    if collected_depots:
+                        if isSteamTools:
+                            await migrate(st_use=True, session=session)
+                            await stool_add(collected_depots, app_id)
+                            log.info('找到SteamTools, 已添加解锁文件')
+
+                        if isGreenLuma:
+                            await migrate(st_use=False, session=session)
+                            await greenluma_add([app_id])
+                            depot_config = {'depots': {depot_id: {
+                                'DecryptionKey': depot_key} for depot_id, depot_key in collected_depots}}
+                            await depotkey_merge(steam_path / 'config' / 'config.vdf', depot_config)
+                            if await greenluma_add([int(i) for i in depot_config['depots'] if i.isdecimal()]):
+                                log.info('找到GreenLuma, 已添加解锁文件')
+
+                        log.info(f'清单最后更新时间: {latest_date}')
+                        log.info(f'入库成功: {app_id}')
+                        os.system('pause')
+                        return True
+
+        log.error(f'清单下载或生成失败: {app_id}')
+        os.system('pause')
+        return False
+
 
 if __name__ == '__main__':
     try:
-        asyncio.run(run())
+        init()
+        if len(sys.argv) > 2:
+            app_id = sys.argv[2]
+        else:
+            app_id = input("App ID: ")
+        asyncio.run(main(app_id))
     except KeyboardInterrupt:
         log.info("程序已退出")
     except SystemExit: