""" 密钥轮换服务 """ 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() # 检查是否有激活的有效密钥(主is_active,辅expire_at) active_key = db.query(APIKey).filter( and_( APIKey.is_active == 1, APIKey.is_deleted == 0, APIKey.expire_at > beijing_now() # 未过期 ) ).order_by(APIKey.key_version.desc()).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.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: """获取当前激活密钥的版本(主is_active,辅expire_at)""" try: active_key = db.query(APIKey).filter( and_( APIKey.is_active == 1, APIKey.is_deleted == 0, APIKey.expire_at > beijing_now() ) ).order_by(APIKey.key_version.desc()).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