Spaces:
Running
Running
File size: 3,387 Bytes
5fb8c4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | """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
|