#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
授权文件生成工具
用于生成和签名授权文件
"""

import json
import os
import sys
import argparse
from datetime import datetime, timedelta
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa, padding
import base64

class LicenseGenerator:
    """授权文件生成器"""
    
    def __init__(self, private_key_path=None):
        """
        初始化授权生成器
        
        Args:
            private_key_path: 私钥文件路径
        """
        self.private_key = None
        if private_key_path and os.path.exists(private_key_path):
            self.load_private_key(private_key_path)
    
    def generate_key_pair(self, private_key_path, public_key_path):
        """
        生成RSA密钥对
        
        Args:
            private_key_path: 私钥保存路径
            public_key_path: 公钥保存路径
        """
        # 生成私钥
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048
        )
        
        # 保存私钥
        with open(private_key_path, 'wb') as f:
            f.write(private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            ))
        
        # 保存公钥
        public_key = private_key.public_key()
        with open(public_key_path, 'wb') as f:
            f.write(public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            ))
        
        self.private_key = private_key
        print(f"密钥对已生成:")
        print(f"  私钥: {private_key_path}")
        print(f"  公钥: {public_key_path}")
    
    def load_private_key(self, private_key_path):
        """
        加载私钥
        
        Args:
            private_key_path: 私钥文件路径
        """
        with open(private_key_path, 'rb') as f:
            self.private_key = serialization.load_pem_private_key(
                f.read(),
                password=None
            )
    
    def generate_license(self, license_info, output_path):
        """
        生成授权文件
        
        Args:
            license_info: 授权信息字典
            output_path: 输出文件路径
        """
        if not self.private_key:
            raise ValueError("未加载私钥，无法生成签名")
        
        # 创建授权数据（与 LicenseManager 要求的顶层结构一致）
        license_data = {
            "product": "BodyBalanceEvaluation",
            "version": license_info.get("version", "1.0"),
            "license_id": license_info.get("license_id"),
            "license_type": license_info.get("license_type"),
            "company_name": license_info.get("company_name", ""),
            "contact_info": license_info.get("contact_info", ""),
            "issued_at": license_info.get("issued_at"),
            "expires_at": license_info.get("expires_at"),
            "machine_id": license_info.get("machine_id", "*"),
            "features": license_info.get("features", {}),
        }

        # 生成签名（对除 signature 外的全部字段进行排序后签名）
        sorted_json = json.dumps(license_data, sort_keys=True, separators=(',', ':'))
        signature = self.private_key.sign(
            sorted_json.encode('utf-8'),
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )

        # 将签名编码为base64并写入
        license_data["signature"] = base64.b64encode(signature).decode('utf-8')

        # 保存授权文件
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(license_data, f, indent=2, ensure_ascii=False)
        
        print(f"授权文件已生成: {output_path}")
        return license_data

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='授权文件生成工具')
    parser.add_argument('--generate-keys', action='store_true', help='生成密钥对')
    parser.add_argument('--private-key', help='私钥文件路径')
    parser.add_argument('--public-key', help='公钥文件路径')
    parser.add_argument('--license-id', help='授权ID')
    parser.add_argument('--license-type', choices=['trial', 'standard', 'professional'], 
                       default='standard', help='授权类型')
    parser.add_argument('--company-name', help='公司名称')
    parser.add_argument('--contact-info', help='联系信息')
    parser.add_argument('--expires-days', type=int, default=365, help='有效期天数')
    parser.add_argument('--machine-id', help='机器ID（留空表示不限制）')
    parser.add_argument('--output', help='输出授权文件路径')
    
    args = parser.parse_args()
    
    generator = LicenseGenerator()
    
    # 生成密钥对
    if args.generate_keys:
        if not args.private_key or not args.public_key:
            print("错误: 生成密钥对需要指定 --private-key 和 --public-key 参数")
            return
        
        generator.generate_key_pair(args.private_key, args.public_key)
        return
    
    # 生成授权文件
    if not args.license_id or not args.company_name or not args.output:
        print("错误: 生成授权文件需要指定 --license-id, --company-name 和 --output 参数")
        return
    
    if not args.private_key or not os.path.exists(args.private_key):
        print("错误: 需要指定有效的私钥文件路径")
        return
    
    # 加载私钥
    generator.load_private_key(args.private_key)
    
    # 创建授权信息
    from datetime import timezone
    now = datetime.now(timezone.utc)
    expires_at = now + timedelta(days=args.expires_days)
    
    # 根据授权类型设置功能
    features = {
        "recording": True,
        "export": True
    }
    
    if args.license_type == 'trial':
        features.update({
            "trial_limit_minutes": 30,
            "max_sessions": 10
        })
    elif args.license_type == 'standard':
        features.update({
            "max_sessions": 1000,
            "max_users": 10
        })
    elif args.license_type == 'professional':
        features.update({
            "max_sessions": -1,  # 无限制
            "max_users": -1,     # 无限制
            "advanced_analytics": True
        })
    
    license_info = {
        "license_id": args.license_id,
        "license_type": args.license_type,
        "company_name": args.company_name,
        "contact_info": args.contact_info or "",
        "issued_at": now.isoformat(),
        "expires_at": expires_at.isoformat(),
        "machine_id": args.machine_id or "*",
        "features": features,
        "version": "1.0"
    }
    
    # 生成授权文件
    generator.generate_license(license_info, args.output)
    
    print("\n授权信息:")
    print(f"  授权ID: {license_info['license_id']}")
    print(f"  授权类型: {license_info['license_type']}")
    print(f"  公司名称: {license_info['company_name']}")
    print(f"  有效期: {license_info['expires_at']}")
    print(f"  功能: {', '.join(license_info['features'].keys())}")

if __name__ == '__main__':
    main()