init_keys.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. """
  2. 初始化密钥对
  3. """
  4. from sqlalchemy.orm import Session
  5. from app.database import SessionLocal
  6. from app.models.api_key import APIKey
  7. from app.core.encrypt import encryptor
  8. from app.core.logger import logger
  9. from datetime import datetime, timedelta
  10. from app.utils.timezone import beijing_now
  11. from app.config import settings
  12. # 初始化密钥过期天数(硬编码)
  13. INIT_KEY_EXPIRE_DAYS = 10
  14. def init_keys():
  15. """初始化密钥对"""
  16. db = SessionLocal()
  17. try:
  18. # 检查是否已有激活的密钥
  19. existing_key = db.query(APIKey).filter(APIKey.is_active == 1).first()
  20. if existing_key:
  21. logger.info("已存在激活的密钥,加载私钥")
  22. encryptor.load_private_key(existing_key.private_key)
  23. return
  24. # 生成新密钥对
  25. logger.info("生成新的RSA密钥对")
  26. private_pem, public_pem = encryptor.generate_key_pair()
  27. # 获取当前最大版本号
  28. max_version = db.query(APIKey).order_by(APIKey.key_version.desc()).first()
  29. next_version = (max_version.key_version + 1) if max_version else 1
  30. # 硬编码初始化密钥过期天数
  31. expire_days = INIT_KEY_EXPIRE_DAYS
  32. # 计算过期时间
  33. expire_at = beijing_now() + timedelta(days=expire_days)
  34. # 创建密钥记录
  35. api_key = APIKey(
  36. public_key=public_pem,
  37. private_key=private_pem,
  38. key_version=next_version,
  39. is_active=1,
  40. created_at=beijing_now(),
  41. expire_at=expire_at
  42. )
  43. db.add(api_key)
  44. db.commit()
  45. logger.info(f"密钥对生成成功,版本号: {next_version}")
  46. except Exception as e:
  47. db.rollback()
  48. logger.exception("初始化密钥失败")
  49. raise
  50. finally:
  51. db.close()