app/services/admin_guard.py

"""Admin 端点 IP 失败计数 + 临时封禁防护。

防止扫描器暴力探测 admin key：连续鉴权失败 N 次后临时封禁该 IP。
内存计数（进程级），重启清零，适合 HF Space 单实例场景。

用法：在鉴权失败时调用 record_failure(ip)，在鉴权前调用 check_blocked(ip)。
"""
from __future__ import annotations

import logging
import time
from collections import defaultdict
from threading import Lock

logger = logging.getLogger(__name__)

# 配置：连续失败 5 次封禁 15 分钟
_FAIL_THRESHOLD = 5
_BAN_SECONDS = 900  # 15 分钟
# 超过此时间无新失败则重置计数（避免长期挂着的计数器）
_RESET_AFTER = 3600

_lock = Lock()
# ip -> {"fails": int, "last_fail_ts": float, "banned_until": float}
_records: dict[str, dict] = defaultdict(dict)


def check_blocked(ip: str) -> bool:
    """检查 IP 是否被临时封禁。返回 True 表示已封禁（应拒绝请求）。"""
    if not ip:
        return False
    with _lock:
        rec = _records.get(ip)
        if not rec:
            return False
        banned_until = rec.get("banned_until", 0)
        if banned_until > time.time():
            return True
        # 封禁已过期，重置
        if banned_until > 0 and banned_until <= time.time():
            rec.clear()
            logger.info("[admin_guard] IP %s ban expired, cleared", ip)
        return False


def record_failure(ip: str) -> None:
    """记录一次鉴权失败。达到阈值后封禁 IP。"""
    if not ip:
        return
    with _lock:
        rec = _records[ip]
        now = time.time()
        # 长时间无失败则重置计数
        last_ts = rec.get("last_fail_ts", 0)
        if last_ts and (now - last_ts > _RESET_AFTER):
            rec.clear()
        rec["fails"] = rec.get("fails", 0) + 1
        rec["last_fail_ts"] = now
        fails = rec["fails"]
        if fails >= _FAIL_THRESHOLD:
            rec["banned_until"] = now + _BAN_SECONDS
            logger.warning(
                "[admin_guard] IP %s banned for %ds after %d failures",
                ip, _BAN_SECONDS, fails,
            )


def record_success(ip: str) -> None:
    """鉴权成功时重置该 IP 的失败计数。"""
    if not ip:
        return
    with _lock:
        rec = _records.get(ip)
        if rec and rec.get("fails", 0) > 0:
            rec.clear()


def get_status() -> dict:
    """返回当前封禁状态快照（用于 admin 面板查看）。"""
    now = time.time()
    with _lock:
        banned = []
        for ip, rec in _records.items():
            bu = rec.get("banned_until", 0)
            if bu > now:
                banned.append({
                    "ip": ip,
                    "fails": rec.get("fails", 0),
                    "banned_until": int(bu),
                    "banned_remaining_sec": int(bu - now),
                })
        return {
            "banned_count": len(banned),
            "banned_ips": banned,
            "fail_threshold": _FAIL_THRESHOLD,
            "ban_seconds": _BAN_SECONDS,
        }


def clear_ban(ip: str) -> bool:
    """手动解除某 IP 的封禁。"""
    with _lock:
        rec = _records.get(ip)
        if rec and rec.get("banned_until", 0) > 0:
            rec.clear()
            logger.info("[admin_guard] IP %s ban manually cleared", ip)
            return True
        return False