Live-streaming/app/admin_auth.py

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import secrets
import time
from datetime import datetime
from pathlib import Path
from typing import Any


SESSION_COOKIE_NAME = "live_admin_session"


def _utc_now_iso() -> str:
    return datetime.utcnow().replace(microsecond=0).isoformat() + "Z"


def _is_local_ip(ip: str) -> bool:
    ip = (ip or "").strip()
    return ip in {"127.0.0.1", "::1", "::ffff:127.0.0.1", "localhost"}


class AdminAuthManager:
    def __init__(self, data_dir: str | Path, logger, *, session_ttl_sec: int = 12 * 60 * 60):
        self.data_dir = Path(data_dir)
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.path = self.data_dir / "admin_auth.json"
        self.audit_path = self.data_dir / "admin_audit.log"
        self.logger = logger
        self.session_ttl_sec = int(session_ttl_sec)
        self._sessions: dict[str, dict[str, Any]] = {}
        self._data = self._default_data()
        self._load()

    def _default_data(self) -> dict[str, Any]:
        return {
            "version": 1,
            "created_at": "",
            "password_hash": "",
            "password_salt": "",
            "password_updated_at": "",
        }

    def _load(self) -> None:
        if not self.path.exists():
            return
        try:
            data = json.loads(self.path.read_text(encoding="utf-8"))
            if isinstance(data, dict):
                self._data.update(data)
        except Exception as exc:
            self.logger.warning(f"[后台认证] 读取认证文件失败: {exc}")

    def _save(self) -> None:
        tmp = self.path.with_suffix(".tmp")
        tmp.write_text(json.dumps(self._data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        tmp.replace(self.path)

    def _hash_password(self, password: str, salt: bytes) -> str:
        digest = hashlib.scrypt(
            password.encode("utf-8"),
            salt=salt,
            n=2 ** 14,
            r=8,
            p=1,
            dklen=64,
        )
        return base64.b64encode(digest).decode("ascii")

    def is_bootstrapped(self) -> bool:
        return bool(self._data.get("password_hash") and self._data.get("password_salt"))

    def can_bootstrap_ip(self, ip: str) -> bool:
        return _is_local_ip(ip)

    def bootstrap(self, password: str) -> None:
        password = str(password or "")
        if len(password) < 8:
            raise ValueError("后台密码至少 8 位")
        salt = secrets.token_bytes(16)
        now = _utc_now_iso()
        self._data = {
            "version": 1,
            "created_at": self._data.get("created_at") or now,
            "password_hash": self._hash_password(password, salt),
            "password_salt": base64.b64encode(salt).decode("ascii"),
            "password_updated_at": now,
        }
        self._sessions.clear()
        self._save()

    def verify_password(self, password: str) -> bool:
        if not self.is_bootstrapped():
            return False
        try:
            salt = base64.b64decode(self._data["password_salt"])
        except Exception:
            return False
        current = self._hash_password(str(password or ""), salt)
        return hmac.compare_digest(current, self._data.get("password_hash", ""))

    def create_session(self, client_ip: str, user_agent: str = "") -> tuple[str, int]:
        self.cleanup_sessions()
        token = secrets.token_urlsafe(32)
        now = int(time.time())
        expires_at = now + self.session_ttl_sec
        self._sessions[token] = {
            "created_at": now,
            "expires_at": expires_at,
            "client_ip": client_ip or "",
            "user_agent": (user_agent or "")[:240],
        }
        return token, expires_at

    def cleanup_sessions(self) -> None:
        now = int(time.time())
        expired = [token for token, session in self._sessions.items() if int(session.get("expires_at", 0)) <= now]
        for token in expired:
            self._sessions.pop(token, None)

    def get_session(self, token: str) -> dict[str, Any] | None:
        self.cleanup_sessions()
        if not token:
            return None
        session = self._sessions.get(token)
        if not session:
            return None
        session["expires_at"] = int(time.time()) + self.session_ttl_sec
        return session

    def destroy_session(self, token: str) -> None:
        if token:
            self._sessions.pop(token, None)

    def write_audit(self, *, action: str, target: str = "", client_ip: str = "", session_id: str = "", detail: str = "") -> None:
        entry = {
            "at": _utc_now_iso(),
            "action": action,
            "target": target,
            "client_ip": client_ip,
            "session_id": session_id[:12],
            "detail": detail,
        }
        line = json.dumps(entry, ensure_ascii=False)
        try:
            with open(self.audit_path, "a", encoding="utf-8") as handle:
                handle.write(line + "\n")
        except Exception as exc:
            self.logger.warning(f"[后台认证] 写审计日志失败: {exc}")
        self.logger.info(f"[审计] action={action} target={target or '-'} ip={client_ip or '-'} detail={detail or '-'}")