buff/backend/app/core/init_keys.py

"""
初始化密钥对
"""
from sqlalchemy.orm import Session
from app.database import SessionLocal
from app.models.api_key import APIKey
from app.core.encrypt import encryptor
from app.core.logger import logger
from datetime import datetime, timedelta
from app.utils.timezone import beijing_now
from app.config import settings

def init_keys():
    """初始化密钥对"""
    db = SessionLocal()
    try:
        # 检查是否已有激活的密钥
        existing_key = db.query(APIKey).filter(APIKey.is_active == 1).first()
        if existing_key:
            logger.info("已存在激活的密钥，加载私钥")
            encryptor.load_private_key(existing_key.private_key)
            return
        
        # 生成新密钥对
        logger.info("生成新的RSA密钥对")
        private_pem, public_pem = encryptor.generate_key_pair()
        
        # 获取当前最大版本号
        max_version = db.query(APIKey).order_by(APIKey.key_version.desc()).first()
        next_version = (max_version.key_version + 1) if max_version else 1
        
        # 计算过期时间
        expire_at = beijing_now() + timedelta(days=settings.KEY_ROTATION_DAYS)
        
        # 创建密钥记录
        api_key = APIKey(
            public_key=public_pem,
            private_key=private_pem,
            key_version=next_version,
            is_active=1,
            created_at=beijing_now(),
            expire_at=expire_at
        )
        db.add(api_key)
        db.commit()
        
        logger.info(f"密钥对生成成功，版本号: {next_version}")
        
    except Exception as e:
        db.rollback()
        logger.exception("初始化密钥失败")
        raise
    finally:
        db.close()