laser_cal/keygen.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RSA 密钥对生成与 License 签发工具

使用方法：
    python keygen.py genkey
    python keygen.py genlicense --expire "never" --product "laser_cal" --version "1.0.0" --licensee "客户名"
"""

import os
import sys
import json
import argparse
import base64
import uuid
from datetime import datetime
from pathlib import Path

try:
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.asymmetric import rsa, padding
    from cryptography.hazmat.backends import default_backend
except ImportError:
    print("错误: 缺少 cryptography 库，请先安装: pip install cryptography")
    sys.exit(1)

BASE_DIR = Path(__file__).parent
LICENSE_DIR = BASE_DIR / "dist" / "license"
PRIVATE_KEY_FILE = LICENSE_DIR / "private_key.pem"
PUBLIC_KEY_FILE = LICENSE_DIR / "public_key.pem"
LICENSE_FILE = LICENSE_DIR / "license.lic"
RSA_KEY_SIZE = 2048


def ensure_dir():
    LICENSE_DIR.mkdir(parents=True, exist_ok=True)


def generate_rsa_keys():
    """生成 RSA 密钥对"""
    ensure_dir()
    private_key = rsa.generate_private_key(
        public_exponent=65537, key_size=RSA_KEY_SIZE, backend=default_backend()
    )
    private_pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption()
    )
    public_pem = private_key.public_key().public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )
    PRIVATE_KEY_FILE.write_bytes(private_pem)
    PUBLIC_KEY_FILE.write_bytes(public_pem)
    print(f"[成功] RSA 密钥对已生成:")
    print(f"  私钥: {PRIVATE_KEY_FILE}")
    print(f"  公钥: {PUBLIC_KEY_FILE}")
    print(f"  密钥长度: {RSA_KEY_SIZE} bits")
    print(f"\n[重要] 请妥善保管私钥文件，切勿泄露！")


def generate_license(expire_date, product="laser_cal", version="1.0.0",
                     licensee="default", max_instances=1):
    """生成 License 文件"""
    ensure_dir()
    if not PRIVATE_KEY_FILE.exists():
        print("[错误] 私钥文件不存在，请先执行 genkey。")
        sys.exit(1)

    private_key = serialization.load_pem_private_key(
        PRIVATE_KEY_FILE.read_bytes(), password=None, backend=default_backend()
    )

    license_data = {
        "license_id": str(uuid.uuid4()),
        "product": product,
        "version": version,
        "licensee": licensee,
        "issue_date": datetime.now().strftime("%Y-%m-%d"),
        "expire_date": expire_date,
        "max_instances": max_instances,
        "features": ["all"]
    }

    license_json = json.dumps(license_data, sort_keys=True, ensure_ascii=False)
    signature = private_key.sign(
        license_json.encode("utf-8"), padding.PKCS1v15(), hashes.SHA256()
    )

    license_content = {
        "license_data": license_data,
        "signature": base64.b64encode(signature).decode("utf-8")
    }
    LICENSE_FILE.write_text(json.dumps(license_content, indent=2, ensure_ascii=False), encoding="utf-8")

    print(f"[成功] License 文件已生成: {LICENSE_FILE}")
    print(f"  License ID : {license_data['license_id']}")
    print(f"  产品       : {license_data['product']}")
    print(f"  版本       : {license_data['version']}")
    print(f"  被授权人   : {license_data['licensee']}")
    print(f"  过期日期   : {license_data['expire_date']}")


def main():
    parser = argparse.ArgumentParser(description="RSA License 工具")
    subparsers = parser.add_subparsers(dest="command")

    subparsers.add_parser("genkey", help="生成 RSA 密钥对")

    lic_parser = subparsers.add_parser("genlicense", help="生成 License")
    lic_parser.add_argument("--expire", type=str, default="never")
    lic_parser.add_argument("--product", type=str, default="laser_cal")
    lic_parser.add_argument("--version", type=str, default="1.0.0")
    lic_parser.add_argument("--licensee", type=str, default="default")
    lic_parser.add_argument("--max-instances", type=int, default=1)

    args = parser.parse_args()

    if args.command == "genkey":
        generate_rsa_keys()
    elif args.command == "genlicense":
        expire = args.expire
        if expire.lower() != "never":
            try:
                datetime.strptime(expire, "%Y-%m-%d")
            except ValueError:
                print("[错误] 日期格式不正确，请使用 YYYY-MM-DD 或 'never'。")
                sys.exit(1)
        generate_license(expire, args.product, args.version, args.licensee, args.max_instances)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()