buff2.0/api-key-tasks/run.py

#!/usr/bin/env python3
"""
密钥过期检查定时任务 - 独立进程
从 backend-admin 拆出，单独运行。需能访问 backend-admin 的 app 包（数据库、配置、模型等）。

运行方式（在项目根目录 /Users/wuzilong/work/00code/steam 下）：
  PYTHONPATH=backend-admin python api-key-tasks/run.py
或先安装依赖后：
  cd api-key-tasks && pip install -r requirements.txt
  PYTHONPATH=../backend-admin python run.py
"""
import asyncio
import signal
import sys
from pathlib import Path

# 将 backend-admin 加入 path，以便 import app.*
_REPO_ROOT = Path(__file__).resolve().parent.parent
_BACKEND_ADMIN = _REPO_ROOT / "backend-admin"
if _BACKEND_ADMIN.exists() and str(_BACKEND_ADMIN) not in sys.path:
    sys.path.insert(0, str(_REPO_ROOT))
    # 以项目根为基准时，需要能 import app -> 即 backend-admin 为当前“应用根”
    sys.path.insert(0, str(_BACKEND_ADMIN))

from datetime import timedelta  # noqa: E402

from app.database import get_db  # noqa: E402
from app.core.logger import logger  # noqa: E402
from app.core.encrypt import encryptor  # noqa: E402
from app.models.api_key import APIKey  # noqa: E402
from app.models.system_config import SystemConfig  # noqa: E402
from app.utils.timezone import beijing_now  # noqa: E402
from app.utils.history import create_api_key_history, create_operation_log  # noqa: E402


_running = True


def _sig_handler(*_):
    global _running
    _running = False


async def key_expiry_check_task():
    """密钥过期检查：当前密钥即将过期时自动生成新密钥（与 backend-admin 原逻辑一致）"""
    check_interval_seconds = SystemConfig.get_config_value("key_expiry_check_interval_seconds", 60)
    advance_generate_minutes = SystemConfig.get_config_value("key_advance_generate_minutes", 10)

    logger.info(
        "api-key-tasks 启动密钥过期检查 (检查间隔: {}秒, 提前生成: {}分钟)",
        check_interval_seconds,
        advance_generate_minutes,
    )

    while _running:
        try:
            logger.debug("执行密钥过期检查...")
            db = next(get_db())
            try:
                current_key = (
                    db.query(APIKey)
                    .filter(APIKey.is_active == 1)
                    .order_by(APIKey.key_version.desc())
                    .first()
                )

                if current_key and current_key.expire_at:
                    beijing_time = beijing_now()
                    if beijing_time.tzinfo is not None:
                        beijing_time_naive = beijing_time.replace(tzinfo=None)
                    else:
                        beijing_time_naive = beijing_time

                    time_until_expire = current_key.expire_at - beijing_time_naive
                    time_until_expire_minutes = time_until_expire.total_seconds() / 60

                    if time_until_expire_minutes <= advance_generate_minutes:
                        existing_new_key = (
                            db.query(APIKey)
                            .filter(
                                APIKey.key_version == current_key.key_version + 1,
                                APIKey.is_active == 1,
                            )
                            .first()
                        )

                        if not existing_new_key:
                            logger.info(
                                "检测到密钥即将过期（剩余{:.1f}分钟），开始自动生成新密钥...",
                                time_until_expire_minutes,
                            )
                            new_key_pair = encryptor.generate_key_pair()
                            new_version = current_key.key_version + 1
                            expire_minutes = SystemConfig.get_config_value("key_expire_minutes", 60)
                            expire_at = beijing_now() + timedelta(minutes=expire_minutes)

                            new_key = APIKey(
                                public_key=new_key_pair[1],
                                private_key=new_key_pair[0],
                                key_version=new_version,
                                is_active=1,
                                expire_at=expire_at,
                                created_at=beijing_now(),
                                updated_at=beijing_now(),
                            )
                            db.add(new_key)
                            db.flush()
                            create_api_key_history(db, new_key, "INSERT")
                            create_operation_log(
                                db=db,
                                admin_id=0,
                                module="密钥管理",
                                action="自动密钥轮换",
                                description=(
                                    f"自动检测密钥即将过期，生成新密钥：版本 {new_version}，"
                                    f"距离过期 {time_until_expire_minutes:.1f}分钟"
                                ),
                                old_data={
                                    "old_version": current_key.key_version,
                                    "expire_at": current_key.expire_at.isoformat(),
                                },
                                new_data={
                                    "new_version": new_version,
                                    "expire_minutes": expire_minutes,
                                },
                                ip_address="api-key-tasks",
                            )
                            db.commit()
                            logger.info(
                                "自动生成新密钥成功：版本 {}，过期时间 {}",
                                new_version,
                                expire_at,
                            )
                        else:
                            logger.debug(
                                "新密钥已存在（版本 {}），无需重复生成",
                                existing_new_key.key_version,
                            )
                    else:
                        logger.debug(
                            "密钥状态正常，距离过期还有 {:.1f} 分钟",
                            time_until_expire_minutes,
                        )
                else:
                    logger.warning("未找到激活的密钥或密钥无过期时间")
            finally:
                db.close()
        except asyncio.CancelledError:
            break
        except Exception as e:
            logger.exception("密钥过期检查任务异常: {}", e)

        # 等待下一轮
        try:
            await asyncio.wait_for(asyncio.shield(asyncio.sleep(check_interval_seconds)), timeout=None)
        except asyncio.CancelledError:
            break

    logger.info("api-key-tasks 密钥过期检查任务已退出")


async def main():
    global _running
    for sig in (signal.SIGINT, signal.SIGTERM):
        try:
            signal.signal(sig, _sig_handler)
        except (ValueError, OSError):
            pass  # Windows / 非主线程等可能不支持

    task = asyncio.create_task(key_expiry_check_task())
    try:
        await task
    except asyncio.CancelledError:
        _running = False
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass


if __name__ == "__main__":
    logger.info("启动 api-key-tasks（密钥定时生成任务）")
    asyncio.run(main())