main.py

import logging
import winreg
import sys
import os
import pyperclip as clip
import typer
from typing import Optional

# 假设这些模块都在同级目录下
from utils import secret_encrypt, write_reg, read_reg
from get_unique_address import secret_get_unique_address
from gold import get_gold_donate_code
from update_key import generate_update_key, analyze_update_key

# 初始化 Typer 应用
app = typer.Typer(add_completion=False, help="PCL 启动器辅助工具")

# 全局变量用于备份恢复
theme_backup = None
gold_backup = None


def setup_logging():
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )


logger = logging.getLogger(__name__)


def load_key():
    # Nuitka 一体化时，数据文件会解压到 sys._MEIPASS
    if hasattr(sys, "_MEIPASS"):
        base_path = sys._MEIPASS
    else:
        base_path = os.path.dirname(__file__)

    key_path = os.path.join(base_path, "key.pem")
    return key_path


def backup_reg():
    """备份注册表值"""
    global theme_backup, gold_backup
    theme_backup = read_reg(
        winreg.HKEY_CURRENT_USER,
        r"Software\PCL",
        "UiLauncherThemeHide2"
    )
    gold_backup = read_reg(
        winreg.HKEY_CURRENT_USER,
        r"Software\PCL",
        "UiLauncherThemeGold"
    )
    logger.info("Registry backed up successfully.")


def restore_reg():
    """出错时恢复注册表"""
    global theme_backup, gold_backup
    write_reg(
        winreg.HKEY_CURRENT_USER,
        r"Software\PCL",
        "UiLauncherThemeHide2",
        theme_backup,
    )
    write_reg(
        winreg.HKEY_CURRENT_USER,
        r"Software\PCL",
        "UiLauncherThemeGold",
        gold_backup
    )
    logger.info("Registry restored.")


def unlock_all_theme(unique_addr: str):
    final_input_key = "PCL" + unique_addr
    origin = "0|1|2|3|4|5|6|7|8|9|10|11|12|13"  # 解锁所有主题
    encrypted_result = secret_encrypt(origin, final_input_key)
    logger.info(
        "successfully got encrypted UiLauncherThemeHide2 value: \n\n%s\n",
        encrypted_result,
    )

    logger.info("writing reg: UiLauncherThemeHide2")
    write_reg(
        winreg.HKEY_CURRENT_USER,
        r"Software\PCL",
        "UiLauncherThemeHide2",
        encrypted_result,
    )
    logger.info(
        "successfully written to HKEY_CURRENT_USER\\Software\\PCL\\UiLauncherThemeHide2"
    )

    gold_code = get_gold_donate_code(unique_addr, key_path=load_key())
    logger.info("encrypting gold donate code")

    gold_code_reg = secret_encrypt(gold_code, final_input_key)
    logger.info("successfully got encrypted gold reg value:\n\n%s\n", gold_code_reg)

    logger.info("writing reg: UiLauncherThemeGold")
    write_reg(
        winreg.HKEY_CURRENT_USER,
        r"Software\PCL",
        "UiLauncherThemeGold",
        gold_code_reg,
    )
    logger.info(
        "successfully written to HKEY_CURRENT_USER\\Software\\PCL\\UiLauncherThemeGold"
    )

    logger.info("done")


# --- CLI Commands ---

@app.command()
def theme():
    """
    默认操作：解锁所有 PCL 主题。
    """
    unique_addr = secret_get_unique_address()
    logger.info("successfully got unique address: %s", unique_addr)

    # 在操作前先备份，否则出错时 restore_reg 没数据可用
    backup_reg()

    try:
        unlock_all_theme(unique_addr)
    except Exception:
        # 一旦出错，恢复之前备份的值
        logger.exception("Fatal error occurred during theme unlock")
        restore_reg()


@app.command()
def update(
        remote_key: Optional[str] =
        typer.Argument(None, help="自定义 remote_update_key（默认为 `linkv3`，在版本 2.11.2 测试可用，如果提示“更新密钥已过期”则可以改高点试试）。")):
    """
    生成更新密钥并自动复制到剪贴板。
    """
    unique_addr = secret_get_unique_address()
    logger.info("successfully got unique address: %s", unique_addr)

    if remote_key:
        generated_key = generate_update_key(unique_addr, load_key(), remote_key)
    else:
        generated_key = generate_update_key(unique_addr, load_key())

    clip.copy(generated_key)
    logger.info(
        f"更新密钥（已自动复制）：\n\n{generated_key}\n\n"
        "在快照版 PCL 中点击【设置 → 其他 → 启动器/系统 → 检查更新】就能使用更新密钥啦！"
    )


@app.command()
def analyze(
        update_key: Optional[str] =
        typer.Argument(help="有效的更新密钥")):
    """
    分析更新密钥，解析出其中的元素。
    """
    unique_addr = secret_get_unique_address()
    logger.info("successfully got unique address: %s", unique_addr)

    analyze_update_key(update_key, unique_addr, load_key())


if __name__ == "__main__":
    setup_logging()

    # 如果不带参数，默认执行 'theme' 命令
    if len(sys.argv) == 1:
        sys.argv.append("theme")

    app()