archery/package.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
应用打包脚本
根据 app.yaml 中列出的文件，打包成 zip 文件
版本号从 version.py 中读取
"""
import argparse
import os
import yaml
import zipfile
from datetime import datetime
import sys
import secrets

MAGIC = b"AROTAE1"  # 7 bytes: Archery OTA Encrypted v1
GCM_NONCE_LEN = 12
GCM_TAG_LEN = 16

# 添加当前目录到路径，以便导入 version 模块
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))


def load_app_yaml(yaml_path='app.yaml'):
    """加载 app.yaml 文件"""
    try:
        with open(yaml_path, 'r', encoding='utf-8') as f:
            return yaml.safe_load(f)
    except Exception as e:
        print(f"[ERROR] 读取 {yaml_path} 失败: {e}")
        return None


def check_files_exist(files, base_dir='.'):
    """检查文件是否存在"""
    missing_files = []
    existing_files = []

    for file_path in files:
        full_path = os.path.join(base_dir, file_path)
        if os.path.exists(full_path):
            existing_files.append(file_path)
        else:
            missing_files.append(file_path)

    return existing_files, missing_files


def get_version_from_version_py():
    """从 version.py 读取版本号"""
    try:
        from version import VERSION
        return VERSION
    except ImportError:
        print("[WARNING] 无法导入 version.py，使用默认版本号 1.0.0")
        return '1.0.0'
    except Exception as e:
        print(f"[WARNING] 读取 version.py 失败: {e}，使用默认版本号 1.0.0")
        return '1.0.0'


def create_zip_package(app_info, files, output_dir='.', base_dir='.'):
    """创建 zip 打包文件"""
    # 生成输出文件名：{name}_v{version}_{timestamp}.zip
    # 版本号从 version.py 读取，而不是从 app.yaml
    app_name = app_info.get('name', 'app')
    version = get_version_from_version_py()  # 从 version.py 读取版本号
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    zip_filename = f"{app_name}_v{version}_{timestamp}.zip"
    zip_path = os.path.join(output_dir, zip_filename)

    print(f"[INFO] 开始打包: {zip_filename}")
    print(f"[INFO] 包含文件数: {len(files)}")

    try:
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for file_path in files:
                full_path = os.path.join(base_dir, file_path)
                # 使用相对路径作为 zip 内的路径
                zipf.write(full_path, file_path)
                print(f"  ✓ {file_path}")

        # 获取文件大小
        file_size = os.path.getsize(zip_path)
        file_size_mb = file_size / (1024 * 1024)

        print(f"\n[SUCCESS] 打包完成!")
        print(f"  文件名: {zip_filename}")
        print(f"  文件大小: {file_size_mb:.2f} MB ({file_size:,} 字节)")
        print(f"  文件路径: {os.path.abspath(zip_path)}")

        return zip_path
    except Exception as e:
        print(f"[ERROR] 打包失败: {e}")
        import traceback
        traceback.print_exc()
        return None


def _validate_key_hex(key_hex: str) -> bytes:
    if not isinstance(key_hex, str):
        raise ValueError("aead key must be hex string")
    key_hex = key_hex.strip().lower()
    if key_hex.startswith("0x"):
        key_hex = key_hex[2:]
    if len(key_hex) != 64:
        raise ValueError("aead key must be 64 hex chars (32 bytes)")
    try:
        key = bytes.fromhex(key_hex)
    except Exception as e:
        raise ValueError(f"invalid hex key: {e}")
    if len(key) != 32:
        raise ValueError("aead key must be 32 bytes")
    return key


def encrypt_zip_aead(zip_path: str, key_hex: str, out_ext: str = ".enc") -> str:
    """
    Encrypt the whole zip file as one blob:
    output format: MAGIC(7) | nonce(12) | ciphertext(N) | tag(16)
    using AES-256-GCM (AEAD).
    """
    # Lazy import: packaging-only dependency
    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    except Exception as e:
        raise RuntimeError(
            "Missing dependency: cryptography. Install with: pip install cryptography. "
            f"Import error: {e}"
        )

    key = _validate_key_hex(key_hex)
    with open(zip_path, "rb") as f:
        plain = f.read()

    nonce = secrets.token_bytes(GCM_NONCE_LEN)
    aesgcm = AESGCM(key)
    ct_and_tag = aesgcm.encrypt(nonce, plain, None)  # ciphertext || tag (16 bytes)

    enc_path = zip_path + out_ext if out_ext else (zip_path + ".enc")
    with open(enc_path, "wb") as f:
        f.write(MAGIC)
        f.write(nonce)
        f.write(ct_and_tag)

    return enc_path


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="打包 app.yaml 文件列表到 zip，并可选进行 AES-256-GCM 加密输出 .enc")
    parser.add_argument("--aead-key-hex", default=None, help="AES-256-GCM key (64 hex chars = 32 bytes). If set, output encrypted file.")
    parser.add_argument("--keep-zip", action="store_true", help="Keep the plaintext zip when encryption is enabled.")
    parser.add_argument("--out-ext", default=".enc", help="Encrypted output extension appended to zip path. Default: .enc (produces *.zip.enc)")
    args = parser.parse_args()

    print("=" * 60)
    print("应用打包脚本")
    print("=" * 60)

    # 1. 加载 app.yaml
    app_info = load_app_yaml('app.yaml')
    if app_info is None:
        return

    # 从 version.py 读取版本号
    version = get_version_from_version_py()
    
    print(f"\n[INFO] 应用信息:")
    print(f"  ID: {app_info.get('id', 'N/A')}")
    print(f"  名称: {app_info.get('name', 'N/A')}")
    print(f"  版本: {version} (来自 version.py)")
    print(f"  作者: {app_info.get('author', 'N/A')}")
    if app_info.get('version') != version:
        print(f"  [注意] app.yaml 中的版本 ({app_info.get('version', 'N/A')}) 与 version.py 不一致")

    # 2. 获取文件列表
    files = app_info.get('files', [])
    if not files:
        print("[ERROR] app.yaml 中没有找到 files 列表")
        return

    print(f"\n[INFO] 文件列表 ({len(files)} 个文件):")

    # 3. 检查文件是否存在
    existing_files, missing_files = check_files_exist(files)

    if missing_files:
        print(f"\n[WARNING] 以下文件不存在，将被跳过:")
        for f in missing_files:
            print(f"  ✗ {f}")

    if not existing_files:
        print("\n[ERROR] 没有找到任何有效文件，无法打包")
        return

    print(f"\n[INFO] 找到 {len(existing_files)} 个有效文件")

    # 4. 创建 zip 包
    zip_path = create_zip_package(app_info, existing_files)

    if zip_path:
        enc_path = None
        if args.aead_key_hex:
            try:
                enc_path = encrypt_zip_aead(zip_path, args.aead_key_hex, out_ext=args.out_ext)
                enc_size = os.path.getsize(enc_path)
                print(f"\n[SUCCESS] AEAD加密完成: {os.path.basename(enc_path)} ({enc_size:,} bytes)")
                print(f"  文件路径: {os.path.abspath(enc_path)}")
                if not args.keep_zip:
                    try:
                        os.remove(zip_path)
                        print(f"[INFO] 已删除明文zip: {os.path.basename(zip_path)}")
                    except Exception as e:
                        print(f"[WARNING] 删除明文zip失败（可忽略）: {e}")
            except Exception as e:
                print(f"\n[ERROR] AEAD加密失败: {e}")
                print("[ERROR] 保留明文zip用于排查。")

        print("\n" + "=" * 60)
        print("打包成功完成!")
        print("=" * 60)
    else:
        print("\n" + "=" * 60)
        print("打包失败!")
        print("=" * 60)


if __name__ == "__main__":
    main()