chongming-core/test/test_build_flags.py

"""
build 命令开关（--encrypt / --license）冒烟测试

两个开关默认都是关的（不加参数 = 明文模块 + 不带授权的内部调试产物），
这里验证"默认关"与"显式打开"在构建产物上都真的生效：
  1. 参数解析：默认 encrypt/license 全关，--encrypt / --license 打开对应开关；
     --only 里 license_guard 是合法步骤（文档里一直写着，之前被漏在校验表外）
  2. 默认（不加密）：.mbank 存明文字节码（encrypted=false）、不写 .key（并清掉旧
     的）、无 packer_key 记录，且**不需要密钥即可导入**；加 --encrypt 那条路仍然
     加密、仍然写 .key，且用 .key 里的密钥能正常导入
  3. 授权构建：prepare 把"是否为授权版"烘焙进 chongming/_build_flags.py（编译进
     产物、部署侧不可改）—— 不加 --license = 非授权版（不校验、插件照常加载），
     加 --license = 授权版（注入并内置公钥、强制校验）；--license 缺公钥默认构建
     失败，显式放行时降级为非授权版（而不是拒绝所有插件的砖头包）；非法公钥报错

运行: uv run python test/test_build_flags.py
"""
import importlib
import os
import shutil
import sys
import tempfile
import tomllib
from pathlib import Path

SRC = Path(__file__).parent
sys.path.insert(0, str(SRC.parent / "utils" / "cli" / "src"))

from module_bank import PythonToSQLite, layout  # noqa: E402
from sqlite_vfs.core import SQLiteVFS  # noqa: E402

from chongming_cli import build as b  # noqa: E402

failures = []


def check(name, cond, detail=""):
    status = "PASS" if cond else "FAIL"
    print(f"[{status}] {name}" + (f" — {detail}" if detail else ""))
    if not cond:
        failures.append(name)


tmp = Path(tempfile.mkdtemp(prefix="build_flags_test_"))

# ========== 1. 参数解析 ==========
_argv, sys.argv = sys.argv, ["build"]
mode, steps, options = b._parse_args()
check(
    "默认参数：pyinstaller + 全流程 + 不加密 + 不授权",
    mode == "pyinstaller" and not steps and not options.encrypt and not options.license,
)

sys.argv = ["build", "--encrypt", "--license", "nuitka"]
mode, steps, options = b._parse_args()
check(
    "--encrypt/--license 打开对应开关",
    mode == "nuitka" and options.encrypt and options.license,
    f"mode={mode} encrypt={options.encrypt} license={options.license}",
)

sys.argv = ["build", "--only", "mbank,license_guard"]
mode, steps, options = b._parse_args()
check("--only 接受 license_guard", steps == {"mbank", "license_guard"}, str(steps))

sys.argv = ["build", "--only", "mbank", "--encrypt", "pyinstaller"]
mode, steps, options = b._parse_args()
check(
    "位置参数与开关可混用",
    mode == "pyinstaller" and steps == {"mbank"} and options.encrypt,
)

sys.argv = ["build", "--unknown-flag"]
try:
    b._parse_args()
    check("未知参数应退出", False)
except SystemExit:
    check("未知参数应退出", True)
sys.argv = _argv


# ========== 2. mbank 加密 / 明文字节码 ==========
def make_source(root: Path) -> Path:
    """造一个最小可导入的包：demo/__init__.py"""
    pkg = root / "demo"
    pkg.mkdir(parents=True)
    (pkg / "__init__.py").write_text("VALUE = 42\n", encoding="utf-8")
    return pkg


def bank_meta(bank: Path, module: str) -> dict:
    vfs = SQLiteVFS(str(bank), compress=False, readonly=True)
    try:
        # 包名 → 包内路径由 module_bank 定，别自己拼；明文未混淆的模块不写
        # meta.json（默认值即"什么都不用记"），用 read_meta 读缺省
        return layout.read_meta(vfs, layout.module_paths(module, True)["meta"])
    finally:
        vfs.close()


def bank_has_key(bank: Path) -> bool:
    vfs = SQLiteVFS(str(bank), compress=False, readonly=True)
    try:
        return vfs.is_file("packer_key")
    finally:
        vfs.close()


def import_from_bank(bank: Path, module: str, key):
    """像 utils/launch.py 那样装载：install_importer → import。返回 (模块, 错误)。"""
    packer = PythonToSQLite(str(bank))
    finder = packer.install_importer(key=key)
    try:
        mod = importlib.import_module(module)
        return mod, None
    except Exception as exc:  # noqa: BLE001 — 测试就是要区分"能导/不能导"
        return None, exc
    finally:
        if finder in sys.meta_path:
            sys.meta_path.remove(finder)
        sys.modules.pop(module, None)


plain_dir = make_source(tmp / "plain_src")
plain_bank = tmp / "plain.mbank"
plain_size = b._pack_mbank(
    str(plain_bank), str(plain_dir), "demo", obfuscate=False, encrypt=False
)
check("明文打包产出 .mbank", plain_bank.exists(), plain_size)
check("明文打包不写 .key", not plain_bank.with_suffix(".key").exists())
check("明文打包 encrypted=false", bank_meta(plain_bank, "demo")["encrypted"] is False)
check("明文打包无 packer_key 记录", not bank_has_key(plain_bank))

mod, err = import_from_bank(plain_bank, "demo", key=None)
check("明文包无密钥即可导入", mod is not None and getattr(mod, "VALUE", None) == 42, str(err))

enc_dir = make_source(tmp / "enc_src")
enc_bank = tmp / "enc.mbank"
enc_key_path = enc_bank.with_suffix(".key")
enc_key_path.write_text("stale-key-from-previous-build", encoding="utf-8")  # 模拟上一次加密构建
enc_size = b._pack_mbank(
    str(enc_bank), str(enc_dir), "demo", obfuscate=False, encrypt=False
)
check("默认（不加密）清掉上一次加密构建留下的 .key", not enc_key_path.exists(), enc_size)

encrypted_bank = tmp / "encrypted.mbank"
key = b.Encryption.generate_key()
_orig_gen = b.Encryption.generate_key
b.Encryption.generate_key = staticmethod(lambda: key)  # 固定密钥，便于用 .key 验证导入
try:
    b._pack_mbank(str(encrypted_bank), str(enc_dir), "demo", obfuscate=False, encrypt=True)
finally:
    b.Encryption.generate_key = _orig_gen

check("--encrypt 打包写 .key", encrypted_bank.with_suffix(".key").exists())
check("--encrypt 打包 encrypted=true", bank_meta(encrypted_bank, "demo")["encrypted"] is True)

mod, err = import_from_bank(encrypted_bank, "demo", key=encrypted_bank.with_suffix(".key").read_text())
check("加密包用 .key 里的密钥可导入", mod is not None, str(err))

mod, err = import_from_bank(encrypted_bank, "demo", key=None)
check("加密包不主动取密钥无法导入（install_importer 不会自己读包内密钥）",
      mod is None and isinstance(err, ImportError), str(err))

# 运行端取密钥的口径：public/utils/launch.py 的 init_module_bank 用 packer.get_key()
# 从**包内**取（它不读旁挂 .key 文件）。打包时没写进包内，加密产物启动就是
# "模块 'plugins' 的字节码是加密的，但没有提供解密密钥" —— 插件全加载失败、
# 契约校验不过，网关直接起不来（实测踩过）
embedded = b.PythonToSQLite(str(encrypted_bank)).get_key()
check("加密包把密钥写进包内（运行端 get_key() 取得到）", embedded == key, f"embedded={embedded!r}")

mod, err = import_from_bank(encrypted_bank, "demo", key=embedded)
check("加密包按运行端口径（get_key() 取出的密钥）可导入", mod is not None, str(err))


# ========== 3. 授权构建：烘焙标记 / 公钥注入 ==========
# 3.0 先验真实模板/骨架的构建契约：下面的用例（含 make_project）都自造 config.toml，
# 所以"模板漏了 [license] 段 / 漏了网关收集项"这类漂移测不出来 —— 而两种漏法都不好在
# 构建期发现（前者首报是"缺少厂商授权公钥"，要补完公钥才暴露"未找到 [license] 段的
# public_key"；后者更隐蔽：打包"成功"，产物启动才报 No module named 'chongming'）。
# 两处契约都在模板（生成新项目用）与根骨架（本仓打包用）的配置里，一起验。
GATEWAY_ENTRY_MODULES = ("chongming", "chongming.app")  # --encrypt 时静态分析看不到
for label, proj_root in (
    ("模板", SRC.parent / "utils" / "template"),
    ("根骨架", SRC.parent),
):
    try:
        injected = b._set_license_pubkey(
            (proj_root / "public" / "config.toml").read_text(encoding="utf-8"), "TESTKEY"
        )
        check(
            f"{label} public/config.toml 可注入 [license] public_key",
            'public_key = "TESTKEY"' in injected,
        )
    except Exception as exc:  # noqa: BLE001 - 缺段/缺键都算契约破坏
        check(f"{label} public/config.toml 可注入 [license] public_key", False, str(exc)[:60])

    # 收集清单要给 --encrypt 兜底：PyArmor 混淆 main.py/server.py 后，打包器跟不到
    # 里面的 `from chongming.app import app`，网关入口包只能靠这里显式声明
    with open(proj_root / "config.toml", "rb") as f:
        pack_list = tomllib.load(f)["pyinstaller_build_import"]
    listed = (
        set(pack_list.get("base", []))
        | set(pack_list.get("linux", []))
        | set(pack_list.get("windows", []))
    )
    missing = [m for m in GATEWAY_ENTRY_MODULES if m not in listed]
    check(
        f"{label} config.toml 收集清单含网关入口包",
        not missing,
        f"缺 {missing}" if missing else "",
    )


def make_project(root: Path) -> Path:
    """造一个最小构建现场：public/{main.py,server.py,config.toml}"""
    (root / "public" / "utils").mkdir(parents=True)
    (root / "public" / "utils" / "launch.py").write_text("", encoding="utf-8")
    (root / "public" / "main.py").write_text("", encoding="utf-8")
    (root / "public" / "server.py").write_text("", encoding="utf-8")
    (root / "public" / "config.toml").write_text(
        '[security]\nsecret_key = "your-secret-key-change-in-production"\n\n'
        '[license]\n# 机器码授权公钥（构建时注入）\npublic_key = ""\n',
        encoding="utf-8",
    )
    return root


def valid_pubkey() -> str:
    """一把合法的 base64 裸 32 字节 Ed25519 公钥（构建期会校验格式）。"""
    import base64

    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

    priv = Ed25519PrivateKey.generate()
    return base64.b64encode(
        priv.public_key().public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw,
        )
    ).decode("ascii")


def flags_path() -> Path:
    """烘焙标记的落盘位置（网关包内；CLI 用同一方式定位）。"""
    spec = importlib.util.find_spec("chongming")
    assert spec and spec.submodule_search_locations, "构建 venv 未安装 chongming-gateway"
    return Path(next(iter(spec.submodule_search_locations))) / "_build_flags.py"


def read_flags() -> str:
    return flags_path().read_text(encoding="utf-8")


env_keys = (b.LICENSE_PUBKEY_ENV, b.LICENSE_PUBKEY_FILE_ENV, b.ALLOW_NO_KEY_ENV)
saved_env = {k: os.environ.pop(k, None) for k in env_keys}  # 环境变量不该影响测试结论

cwd = Path.cwd()
proj = make_project(tmp / "proj_default")
flags_file = flags_path()
flags_backup = flags_file.read_text(encoding="utf-8") if flags_file.exists() else None
os.chdir(proj)
try:
    # 默认（不授权）：产物烘焙"非授权版"——不校验授权、插件照常加载（不是砖头包）
    msg = b._stage_prepare_build_dir(None, "pyinstaller", b.BuildOptions())
    conf = (proj / "build" / "config.toml").read_text(encoding="utf-8")
    check("默认（不授权）的 prepare 不报错", msg == "资源准备完成", msg)
    check("默认不注入公钥（模板 public_key 保持空）", 'public_key = ""' in conf)
    check("JWT 密钥仍会注入", "your-secret-key-change-in-production" not in conf)
    check("默认烘焙非授权版标记（enforce=False）", flags_file.exists() and "LICENSE_ENFORCE = False" in read_flags())

    # --license 缺公钥：默认让构建失败（避免打出名不副实的"授权版"）
    try:
        b._stage_prepare_build_dir(None, "pyinstaller", b.BuildOptions(license=True))
        check("--license 缺授权公钥应失败", False)
    except RuntimeError as exc:
        check("--license 缺授权公钥应失败", "缺少厂商授权公钥" in str(exc))

    # --license 缺公钥 + 显式放行：烘焙"非授权版"（旧行为是产出拒绝所有插件的砖头包）
    os.environ[b.ALLOW_NO_KEY_ENV] = "1"
    try:
        b._stage_prepare_build_dir(None, "pyinstaller", b.BuildOptions(license=True))
        check("显式放行时不报错", True)
        check("显式放行时烘焙非授权版标记（enforce=False）",
              "LICENSE_ENFORCE = False" in read_flags())
    finally:
        os.environ.pop(b.ALLOW_NO_KEY_ENV, None)

    # --license + 公钥：注入产物配置 + 烘焙授权版标记（公钥随代码内置）
    pubkey = valid_pubkey()
    os.environ[b.LICENSE_PUBKEY_ENV] = pubkey
    try:
        b._stage_prepare_build_dir(None, "pyinstaller", b.BuildOptions(license=True))
        conf = (proj / "build" / "config.toml").read_text(encoding="utf-8")
        flags_src = read_flags()
        ns: dict = {}
        exec(compile(flags_src, str(flags_file), "exec"), ns)  # 生成物必须是合法 Python
        check("--license + 公钥：注入 [license] public_key", f'public_key = "{pubkey}"' in conf)
        check("--license + 公钥：烘焙授权版标记", ns.get("LICENSE_ENFORCE") is True)
        check("--license + 公钥：公钥随代码内置", ns.get("LICENSE_PUBLIC_KEY") == pubkey)
        check("--license + 公钥：记录公钥指纹", bool(ns.get("PUBLIC_KEY_FINGERPRINT")))
    finally:
        os.environ.pop(b.LICENSE_PUBKEY_ENV, None)

    # 非法公钥：构建期即报错（否则会打出"谁都验不过"的授权版产物）
    os.environ[b.LICENSE_PUBKEY_ENV] = "not-a-real-key"
    try:
        b._stage_prepare_build_dir(None, "pyinstaller", b.BuildOptions(license=True))
        check("非法公钥应终止构建", False)
    except Exception as exc:  # noqa: BLE001 - 具体异常类型由公钥库决定
        check("非法公钥应终止构建", "公钥" in str(exc) or "key" in str(exc).lower(), str(exc)[:60])
    finally:
        os.environ.pop(b.LICENSE_PUBKEY_ENV, None)
finally:
    os.chdir(cwd)
    for k, v in saved_env.items():
        if v is not None:
            os.environ[k] = v
    # 清理写进构建 venv 的烘焙标记（测试不应留下副作用）
    if flags_backup is None:
        flags_file.unlink(missing_ok=True)
    else:
        flags_file.write_text(flags_backup, encoding="utf-8")

shutil.rmtree(tmp, ignore_errors=True)

print("\n" + ("=" * 40))
if failures:
    print(f"❌ {len(failures)} 项未通过: {failures}")
    sys.exit(1)
print("✅ 全部通过")