# -*- coding: utf-8 -*-
# Copyright (c) 2025 relakkes@gmail.com
#
# This file is part of MediaCrawler project.
# Repository: https://github.com/NanmiCoder/MediaCrawler/blob/main/api/auth.py
# GitHub: https://github.com/NanmiCoder
# Licensed under NON-COMMERCIAL LEARNING LICENSE 1.1
#
# 声明：本代码仅供学习和研究目的使用。使用者应遵守以下原则：
# 1. 不得用于任何商业用途。
# 2. 使用时应遵守目标平台的使用条款和robots.txt规则。
# 3. 不得进行大规模爬取或对平台造成运营干扰。
# 4. 应合理控制请求频率，避免给目标平台带来不必要的负担。
# 5. 不得用于任何非法或不当的用途。
#
# 详细许可条款请参阅项目根目录下的LICENSE文件。
# 使用本代码即表示您同意遵守上述原则和LICENSE中的所有条款。

"""Authentication for the WebUI.

Design constraints that drove this, all verified against the codebase:

* **Cookies, not bearer headers, are the primary transport.** Browser
  WebSockets cannot set custom headers on the handshake, and the data-export
  downloads use ``window.open`` (a navigation, also header-less). Only a cookie
  is carried on both. The same opaque token is *also* accepted from an
  ``Authorization: Bearer`` header so ``curl`` and scripts remain usable.
* **Enforcement is a ``Depends``, not middleware.** ``BaseHTTPMiddleware``
  returns early for any non-``http`` scope, so it never sees a WebSocket --
  a middleware-only gate would leave the live log stream wide open. It is also
  overridable per-test via ``app.dependency_overrides``.
* **Sessions are server-side** so logout and password-change revoke immediately.

Only the environment variable ``MC_PASSWORD`` can bypass the stored hash. That is
the documented way back in if the password is forgotten, which is why it is
never persisted.
"""

import asyncio
import base64
import binascii
import hashlib
import hmac
import os
import secrets
import time
from typing import Optional

from anyio import to_thread
from fastapi import HTTPException, Request, WebSocket, WebSocketException, status
from sqlalchemy import delete
from sqlalchemy.ext.asyncio import AsyncSession

from tools.time_util import get_current_timestamp

from .monitor.db import get_session
from .monitor.models import (
    SETTING_AUTH_PASSWORD_HASH,
    SETTING_AUTH_PASSWORD_UPDATED_AT,
    AuthSession,
)
from .monitor.settings import get_setting, set_setting

SESSION_COOKIE_NAME = "mc_session"

# OWASP's current PBKDF2-HMAC-SHA256 guidance. Deliberately slow -- see
# verify_password() for why that cost must not land on the event loop.
PBKDF2_ITERATIONS = 600_000
PBKDF2_ALGO = "pbkdf2_sha256"

# A single generic message for every failure mode, so the response never
# reveals whether a password is set, wrong, or empty.
INVALID_CREDENTIALS = "用户名或密码错误"

# Brute-force throttle. In-process is sufficient: this is a single-user tool and
# uvicorn runs one worker. Documented as reset-on-restart.
THROTTLE_THRESHOLD = 5
THROTTLE_WINDOW_SECONDS = 900
THROTTLE_MAX_LOCKOUT_SECONDS = 900

_failures: dict[str, list[float]] = {}
_throttle_lock = asyncio.Lock()


def _now() -> float:
    """Monotonic clock, indirected so tests can drive it without sleeping."""
    return time.monotonic()


# ---------------------------------------------------------------------------
# Environment configuration (read at call time so tests can set it per-case)
# ---------------------------------------------------------------------------

def env_password() -> str:
    return os.getenv("MC_PASSWORD", "").strip()


def cookie_secure() -> bool:
    return os.getenv("MC_COOKIE_SECURE", "").strip().lower() in ("1", "true", "yes", "y")


def session_ttl_ms() -> int:
    try:
        hours = int(os.getenv("MC_SESSION_TTL_HOURS", "336"))
    except ValueError:
        hours = 336
    return max(hours, 1) * 3_600_000


# ---------------------------------------------------------------------------
# Password hashing
# ---------------------------------------------------------------------------

def _b64(raw: bytes) -> str:
    return base64.b64encode(raw).decode("ascii")


def hash_password(password: str, *, iterations: Optional[int] = None) -> str:
    """Return a self-describing hash so the iteration count can be raised later
    without a migration: ``pbkdf2_sha256$<iterations>$<salt>$<hash>``.

    ``iterations`` is resolved at call time (not bound as a default) so tests can
    lower it; the production value stays the module constant.
    """
    iterations = iterations or PBKDF2_ITERATIONS
    salt = secrets.token_bytes(16)
    digest = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, iterations)
    return f"{PBKDF2_ALGO}${iterations}${_b64(salt)}${_b64(digest)}"


def _verify_password_sync(password: str, stored: str) -> bool:
    try:
        algo, iterations_raw, salt_raw, digest_raw = stored.split("$")
        if algo != PBKDF2_ALGO:
            return False
        salt = base64.b64decode(salt_raw)
        expected = base64.b64decode(digest_raw)
        actual = hashlib.pbkdf2_hmac("sha256", password.encode("utf-8"), salt, int(iterations_raw))
    except (ValueError, TypeError, binascii.Error):
        return False
    return hmac.compare_digest(actual, expected)


async def verify_password(password: str, stored: str) -> bool:
    """Verify off the event loop.

    At 600k iterations this takes a few hundred milliseconds. Running it inline
    in an async handler would block the loop entirely -- stalling the monitor
    scheduler and every websocket ping -- and present as "the whole UI freezes
    when I click login".
    """
    return await to_thread.run_sync(_verify_password_sync, password, stored)


async def current_password_hash(session: AsyncSession) -> str:
    return (await get_setting(session, SETTING_AUTH_PASSWORD_HASH)) or ""


async def set_password(session: AsyncSession, password: str) -> None:
    await set_setting(session, SETTING_AUTH_PASSWORD_HASH, hash_password(password))
    await set_setting(
        session, SETTING_AUTH_PASSWORD_UPDATED_AT, str(get_current_timestamp())
    )


async def check_password(session: AsyncSession, password: str) -> bool:
    """The environment override wins over the stored hash, always.

    That is the escape hatch: forgetting the password is recoverable by setting
    MC_PASSWORD and restarting, without touching the database.
    """
    override = env_password()
    if override:
        return hmac.compare_digest(password, override)

    stored = await current_password_hash(session)
    if not stored:
        return False
    return await verify_password(password, stored)


async def ensure_initial_credential() -> Optional[str]:
    """Seed a password on first run; returns it once so main() can print it.

    Deliberately NOT an unauthenticated "set your password" endpoint: on a
    LAN-exposed bind that is a claim-the-instance race where whoever reaches the
    page first becomes the administrator. Generating and printing a random
    password avoids the race and also avoids locking the operator out.
    """
    if env_password():
        return None

    async with get_session() as session:
        if await current_password_hash(session):
            return None
        generated = secrets.token_urlsafe(12)
        await set_password(session, generated)
        return generated


# ---------------------------------------------------------------------------
# Sessions
# ---------------------------------------------------------------------------

def _hash_token(token: str) -> str:
    return hashlib.sha256(token.encode("utf-8")).hexdigest()


async def create_session(session: AsyncSession) -> tuple[str, int]:
    """Issue a session. Returns (token, expires_at_ms).

    The caller receives the raw token; only its hash is stored.
    """
    token = secrets.token_urlsafe(32)
    now = get_current_timestamp()
    expires_at = now + session_ttl_ms()
    session.add(
        AuthSession(
            token_hash=_hash_token(token),
            created_at=now,
            expires_at=expires_at,
            last_seen_at=now,
        )
    )
    return token, expires_at


async def resolve_session(session: AsyncSession, token: str) -> Optional[AuthSession]:
    if not token:
        return None

    row = await session.get(AuthSession, _hash_token(token))
    if row is None:
        return None

    now = get_current_timestamp()
    if row.expires_at <= now:
        await session.delete(row)
        return None

    row.last_seen_at = now
    return row


async def revoke_session(session: AsyncSession, token: str) -> None:
    row = await session.get(AuthSession, _hash_token(token))
    if row is not None:
        await session.delete(row)


async def revoke_all_sessions(session: AsyncSession) -> int:
    """Used on password change, which is what makes "all devices logged out"
    take effect immediately rather than at token expiry."""
    result = await session.execute(delete(AuthSession))
    return result.rowcount or 0


async def purge_expired_sessions(session: AsyncSession) -> None:
    await session.execute(delete(AuthSession).where(AuthSession.expires_at <= get_current_timestamp()))


# ---------------------------------------------------------------------------
# Credential extraction and enforcement
# ---------------------------------------------------------------------------

def token_from_request(request: Request) -> str:
    """Cookie first (browsers, websockets, navigations), then Bearer (scripts)."""
    token = request.cookies.get(SESSION_COOKIE_NAME, "")
    if token:
        return token
    header = request.headers.get("authorization", "")
    if header.lower().startswith("bearer "):
        return header[7:].strip()
    return ""


def _unauthorized() -> HTTPException:
    return HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail=INVALID_CREDENTIALS,
        headers={"WWW-Authenticate": "Bearer"},
    )


async def require_auth(request: Request) -> None:
    """FastAPI dependency guarding the protected routers.

    Applied per-router via ``include_router(..., dependencies=[Depends(...)])``
    rather than as app-wide middleware, so it appears in the OpenAPI schema,
    returns a correct 401, and can be overridden in tests.
    """
    token = token_from_request(request)
    if not token:
        raise _unauthorized()

    async with get_session() as session:
        if await resolve_session(session, token) is None:
            raise _unauthorized()


async def require_ws_auth(websocket: WebSocket) -> None:
    """Guard for WebSocket routes.

    These need their own dependency: ``BaseHTTPMiddleware`` passes any non-http
    scope straight through, and router-level HTTP dependencies do not apply to
    websocket routes. Raising ``WebSocketException`` closes the handshake with
    the given code; ``HTTPException`` would be meaningless here.
    """
    token = websocket.cookies.get(SESSION_COOKIE_NAME, "")
    if not token:
        raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)

    async with get_session() as session:
        if await resolve_session(session, token) is None:
            raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)


# ---------------------------------------------------------------------------
# Brute-force throttle
# ---------------------------------------------------------------------------

def client_key(request: Request) -> str:
    """Identify the caller for throttling.

    ``X-Forwarded-For`` is only consulted when the operator explicitly opts in,
    because otherwise any client could spoof the header and throttle someone
    else (or evade its own throttle).
    """
    if os.getenv("MC_TRUST_PROXY", "").strip() == "1":
        forwarded = request.headers.get("x-forwarded-for", "")
        if forwarded:
            return forwarded.split(",")[0].strip()
    return request.client.host if request.client else "unknown"


def _recent_failures(key: str) -> list[float]:
    cutoff = _now() - THROTTLE_WINDOW_SECONDS
    return [ts for ts in _failures.get(key, []) if ts >= cutoff]


async def retry_after_seconds(key: str) -> int:
    """0 when not throttled, otherwise how long the caller must wait."""
    async with _throttle_lock:
        recent = _recent_failures(key)
        _failures[key] = recent
        if len(recent) < THROTTLE_THRESHOLD:
            return 0

        # Lockout doubles per failure past the threshold, capped.
        extra = len(recent) - THROTTLE_THRESHOLD
        lockout = min(2 ** extra, THROTTLE_MAX_LOCKOUT_SECONDS)
        elapsed = _now() - recent[-1]
        remaining = int(lockout - elapsed)
        return max(remaining, 1)


async def record_failure(key: str) -> None:
    async with _throttle_lock:
        _failures.setdefault(key, []).append(_now())


async def clear_failures(key: str) -> None:
    async with _throttle_lock:
        _failures.pop(key, None)


def reset_throttle_state() -> None:
    """Test hook: drop all throttle state."""
    _failures.clear()