"""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