personal-progress-monitor/app/auth.py

from __future__ import annotations

import hashlib
import hmac
import os
import secrets
import time
from dataclasses import dataclass, field


SESSION_TTL_SECONDS = 60 * 60 * 24 * 30  # 与登录 Cookie 的 max_age 一致
LOGIN_MAX_FAILURES = 5
LOGIN_LOCKOUT_SECONDS = 60


def make_password_hash(password: str, iterations: int = 310_000) -> str:
    salt = secrets.token_hex(16)
    digest = hashlib.pbkdf2_hmac("sha256", password.encode(), bytes.fromhex(salt), iterations).hex()
    return f"pbkdf2_sha256${iterations}${salt}${digest}"


def verify_password(password: str, encoded: str) -> bool:
    try:
        algorithm, iterations_text, salt, expected = encoded.split("$", 3)
        if algorithm != "pbkdf2_sha256":
            return False
        digest = hashlib.pbkdf2_hmac(
            "sha256", password.encode(), bytes.fromhex(salt), int(iterations_text)
        ).hex()
        return hmac.compare_digest(digest, expected)
    except (ValueError, TypeError):
        return False


@dataclass
class Session:
    csrf: str
    created_at: float = field(default_factory=time.time)


class SessionManager:
    def __init__(
        self,
        password_hash: str | None,
        ttl_seconds: float = SESSION_TTL_SECONDS,
        max_failures: int = LOGIN_MAX_FAILURES,
        lockout_seconds: float = LOGIN_LOCKOUT_SECONDS,
    ):
        self.password_hash = password_hash
        self.ttl_seconds = ttl_seconds
        self.max_failures = max_failures
        self.lockout_seconds = lockout_seconds
        self.sessions: dict[str, Session] = {}
        self._failures: dict[str, list[float]] = {}

    @property
    def enabled(self) -> bool:
        return bool(self.password_hash)

    @classmethod
    def from_environment(cls) -> "SessionManager":
        encoded = os.getenv("PPM_PASSWORD_HASH", "").strip()
        plain = os.getenv("PPM_PASSWORD", "")
        allow_no_auth = os.getenv("PPM_ALLOW_NO_AUTH", "").strip().lower() in {"1", "true", "yes"}

        if not encoded and plain:
            encoded = make_password_hash(plain)

        if not encoded:
            # 默认 fail-secure：没有密码就拒绝启动，除非显式豁免（仅本地开发）。
            if allow_no_auth:
                return cls(None)
            raise RuntimeError(
                "未配置访问密码，拒绝启动（fail-secure）。"
                "请设置环境变量 PPM_PASSWORD 或 PPM_PASSWORD_HASH；"
                "仅本地开发可设置 PPM_ALLOW_NO_AUTH=1 豁免认证。"
            )
        return cls(encoded)

    def _expire(self) -> None:
        cutoff = time.time() - self.ttl_seconds
        expired = [sid for sid, session in self.sessions.items() if session.created_at < cutoff]
        for sid in expired:
            self.sessions.pop(sid, None)

    def _prune_failures(self, now: float) -> None:
        cutoff = now - self.lockout_seconds
        for key in list(self._failures):
            self._failures[key] = [ts for ts in self._failures[key] if ts > cutoff]
            if not self._failures[key]:
                self._failures.pop(key, None)

    def is_login_locked(self, client_key: str) -> bool:
        self._prune_failures(time.time())
        return len(self._failures.get(client_key, [])) >= self.max_failures

    def login(self, password: str, client_key: str = "") -> tuple[str, Session] | None:
        if not self.password_hash:
            return None
        # 锁定检查内聚到 login：即便绕过 HTTP 端点直接调用也无法暴力破解。
        if client_key and self.is_login_locked(client_key):
            return None
        if not verify_password(password, self.password_hash):
            if client_key:
                self._failures.setdefault(client_key, []).append(time.time())
            return None
        self._expire()
        self._failures.pop(client_key, None)
        session_id = secrets.token_urlsafe(32)
        session = Session(csrf=secrets.token_urlsafe(32))
        self.sessions[session_id] = session
        return session_id, session

    def get(self, session_id: str | None) -> Session | None:
        if not session_id:
            return None
        self._expire()
        return self.sessions.get(session_id)

    def logout(self, session_id: str | None) -> None:
        if session_id:
            self.sessions.pop(session_id, None)