102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
"""认证、配置加载与登录限流工具（仅依赖标准库）。"""
|
|
import base64
|
|
import hashlib
|
|
import hmac
|
|
import json
|
|
import os
|
|
import secrets
|
|
import threading
|
|
import time
|
|
|
|
DEFAULT_CONFIG_PATH = os.environ.get("LZWLAB_TRANSFER_CONFIG", "/etc/lzwlab-transfer/config.json")
|
|
|
|
_SCRYPT_N = 2 ** 14
|
|
_SCRYPT_R = 8
|
|
_SCRYPT_P = 1
|
|
|
|
|
|
def _b64encode(raw: bytes) -> str:
|
|
return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")
|
|
|
|
|
|
def _b64decode(text: str) -> bytes:
|
|
return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))
|
|
|
|
|
|
def hash_password(password: str) -> str:
|
|
"""返回 scrypt 密码哈希，格式：scrypt$N$r$p$salt_b64$hash_b64"""
|
|
if not isinstance(password, str) or not password:
|
|
raise ValueError("密码不能为空")
|
|
salt = secrets.token_bytes(16)
|
|
digest = hashlib.scrypt(
|
|
password.encode("utf-8"), salt=salt, n=_SCRYPT_N, r=_SCRYPT_R, p=_SCRYPT_P, dklen=32
|
|
)
|
|
return "scrypt$%d$%d$%d$%s$%s" % (
|
|
_SCRYPT_N, _SCRYPT_R, _SCRYPT_P, _b64encode(salt), _b64encode(digest)
|
|
)
|
|
|
|
|
|
def verify_password(password: str, stored: str) -> bool:
|
|
try:
|
|
parts = stored.split("$")
|
|
if len(parts) != 6 or parts[0] != "scrypt":
|
|
return False
|
|
_, n_s, r_s, p_s, salt_s, digest_s = parts
|
|
n, r, p = int(n_s), int(r_s), int(p_s)
|
|
digest = hashlib.scrypt(
|
|
password.encode("utf-8"), salt=_b64decode(salt_s), n=n, r=r, p=p, dklen=32
|
|
)
|
|
return hmac.compare_digest(digest, _b64decode(digest_s))
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def load_config(path: str | None = None) -> dict:
|
|
path = path or DEFAULT_CONFIG_PATH
|
|
if not os.path.exists(path):
|
|
raise RuntimeError("配置文件不存在：%s" % path)
|
|
with open(path, "r", encoding="utf-8") as fh:
|
|
cfg = json.load(fh)
|
|
if not isinstance(cfg.get("users"), dict) or not cfg["users"]:
|
|
raise RuntimeError("配置缺少 users")
|
|
if not cfg.get("secret_key") or len(str(cfg["secret_key"])) < 32:
|
|
raise RuntimeError("配置缺少 secret_key（至少 32 字符）")
|
|
data_dir = cfg.get("data_dir", "/var/lib/lzwlab-transfer")
|
|
cfg["files_dir"] = os.path.join(data_dir, "files")
|
|
cfg["database_path"] = cfg.get("database_path", os.path.join(data_dir, "files.db"))
|
|
cfg["max_content_length"] = int(cfg.get("max_content_length", 5 * 1024 ** 3))
|
|
cfg["cookie_secure"] = bool(cfg.get("cookie_secure", True))
|
|
cfg["session_days"] = int(cfg.get("session_days", 7))
|
|
return cfg
|
|
|
|
|
|
def random_secret() -> str:
|
|
return secrets.token_urlsafe(48)
|
|
|
|
|
|
class RateLimiter:
|
|
"""基于内存的按 IP 登录失败限流。"""
|
|
|
|
def __init__(self, max_failures: int = 10, window_seconds: int = 900):
|
|
self.max_failures = max_failures
|
|
self.window_seconds = window_seconds
|
|
self._lock = threading.Lock()
|
|
self._failures: dict[str, list[float]] = {}
|
|
|
|
def _prune(self, ip: str, now: float) -> None:
|
|
self._failures[ip] = [t for t in self._failures.get(ip, []) if now - t < self.window_seconds]
|
|
|
|
def is_blocked(self, ip: str) -> bool:
|
|
with self._lock:
|
|
self._prune(ip, time.monotonic())
|
|
return len(self._failures.get(ip, [])) >= self.max_failures
|
|
|
|
def record_failure(self, ip: str) -> None:
|
|
with self._lock:
|
|
self._failures.setdefault(ip, []).append(time.monotonic())
|
|
self._prune(ip, time.monotonic())
|
|
|
|
def reset(self, ip: str) -> None:
|
|
with self._lock:
|
|
self._failures.pop(ip, None)
|