buff/backend/app/services/key_rotation_service.py

"""
密钥轮换服务
"""
import time
from datetime import datetime, timedelta
from sqlalchemy.orm import Session
from sqlalchemy import and_, or_
from app.database import get_db
from app.models.api_key import APIKey
from app.models.system_config import SystemConfig
from app.core.encrypt import encryptor
from app.core.logger import logger
from app.utils.timezone import beijing_now

class KeyRotationService:
    """密钥轮换服务"""

    @staticmethod
    def get_config_dict() -> dict:
        """获取密钥轮换相关配置"""
        return {
            'key_overlap_minutes': SystemConfig.get_config_value('key_overlap_minutes', 30),
            'key_expire_minutes': SystemConfig.get_config_value('key_expire_minutes', 60),
        }

    @staticmethod
    def should_rotate_key(db: Session) -> bool:
        """
        检查是否需要轮换密钥

        轮换条件：
        1. 当前没有激活的密钥
        2. 当前密钥即将过期（剩余时间少于重叠时间）
        """
        try:
            config = KeyRotationService.get_config_dict()

            # 检查是否有激活的密钥
            active_key = db.query(APIKey).filter(
                and_(
                    APIKey.is_active == 1,
                    APIKey.is_deleted == 0,
                    or_(
                        APIKey.expire_at.is_(None),
                        APIKey.expire_at > beijing_now()
                    )
                )
            ).first()

            if not active_key:
                logger.info("没有激活的密钥，需要生成新密钥")
                return True

            # 检查是否即将过期
            if active_key.expire_at:
                current_time = beijing_now()
                # 确保两个时间都是相同类型
                if active_key.expire_at.tzinfo is None:
                    # 如果数据库时间是naive的，转换为aware
                    from datetime import timezone, timedelta
                    expire_time = active_key.expire_at.replace(tzinfo=timezone(timedelta(seconds=28800)))  # 北京时间
                else:
                    expire_time = active_key.expire_at

                time_until_expire = expire_time - current_time
                overlap_minutes = config.get('key_overlap_minutes', 30)

                if time_until_expire.total_seconds() < (overlap_minutes * 60):
                    logger.info(f"当前密钥即将过期（剩余{time_until_expire.total_seconds()/60:.1f}分钟），需要轮换")
                    return True

            return False

        except Exception as e:
            logger.exception("检查密钥轮换条件失败")
            return True  # 出错时认为需要轮换

    @staticmethod
    def rotate_key(db: Session) -> tuple[bool, str, str]:
        """
        执行密钥轮换

        Returns:
            tuple: (是否成功, 新公钥, 错误信息)
        """
        try:
            config = KeyRotationService.get_config_dict()

            # 生成新密钥对
            private_pem, public_pem = encryptor.generate_key_pair()

            # 获取下一个版本号
            max_version = db.query(APIKey.key_version).filter(APIKey.is_deleted == 0).order_by(APIKey.key_version.desc()).first()
            next_version = (max_version[0] + 1) if max_version else 1

            # 计算过期时间
            expire_minutes = config.get('key_expire_minutes', 60)
            expire_at = beijing_now() + timedelta(minutes=expire_minutes)

            # 保存新密钥
            new_key = APIKey(
                public_key=public_pem,
                private_key=private_pem,
                key_version=next_version,
                is_active=1,
                expire_at=expire_at,
                created_at=beijing_now()
            )

            # 停用旧密钥（如果存在）
            db.query(APIKey).filter(
                and_(
                    APIKey.is_active == 1,
                    APIKey.is_deleted == 0
                )
            ).update({
                "is_active": 0,
                "updated_at": beijing_now()
            })

            # 保存新密钥
            db.add(new_key)
            db.commit()
            db.refresh(new_key)

            # 加载新密钥到加密器
            encryptor.load_private_key(private_pem, next_version)

            logger.info(f"密钥轮换成功，新版本: {next_version}")
            return True, public_pem, ""

        except Exception as e:
            db.rollback()
            error_msg = f"密钥轮换失败: {str(e)}"
            logger.exception("密钥轮换失败")
            return False, "", error_msg

    @staticmethod
    def cleanup_expired_keys(db: Session) -> int:
        """
        清理过期的密钥

        Returns:
            int: 清理的密钥数量
        """
        try:
            # 删除过期超过 overlap_minutes 的密钥
            config = KeyRotationService.get_config_dict()
            overlap_minutes = config.get('key_overlap_minutes', 30)
            cutoff_time = beijing_now() - timedelta(minutes=overlap_minutes)

            result = db.query(APIKey).filter(
                and_(
                    APIKey.is_active == 0,  # 只删除非激活的密钥
                    APIKey.is_deleted == 0,
                    APIKey.expire_at < cutoff_time
                )
            ).update({
                "is_deleted": 1,
                "updated_at": beijing_now()
            })

            if result > 0:
                db.commit()
                logger.info(f"清理了 {result} 个过期的密钥")

            return result

        except Exception as e:
            db.rollback()
            logger.exception("清理过期密钥失败")
            return 0

    @staticmethod
    def get_active_key_version(db: Session) -> int:
        """获取当前激活密钥的版本"""
        try:
            active_key = db.query(APIKey).filter(
                and_(
                    APIKey.is_active == 1,
                    APIKey.is_deleted == 0
                )
            ).first()

            return active_key.key_version if active_key else 0

        except Exception as e:
            logger.exception("获取激活密钥版本失败")
            return 0

    @staticmethod
    def get_public_key_by_version(db: Session, version: int) -> str:
        """根据版本获取公钥"""
        try:
            api_key = db.query(APIKey).filter(
                and_(
                    APIKey.key_version == version,
                    APIKey.is_deleted == 0
                )
            ).first()

            return api_key.public_key if api_key else ""

        except Exception as e:
            logger.exception(f"获取版本 {version} 的公钥失败")
            return ""

    @staticmethod
    def initialize_keys(db: Session) -> bool:
        """初始化密钥系统"""
        try:
            # 检查是否已有激活密钥
            if not KeyRotationService.should_rotate_key(db):
                logger.info("密钥系统已初始化")
                return True

            # 执行密钥轮换
            success, _, error_msg = KeyRotationService.rotate_key(db)
            if success:
                logger.info("密钥系统初始化成功")
                return True
            else:
                logger.error(f"密钥系统初始化失败: {error_msg}")
                return False

        except Exception as e:
            logger.exception("密钥系统初始化异常")
            return False