WeReadBoard/auth.py

import hmac
import os
import secrets
from datetime import timedelta

import pymysql
from flask import Blueprint, current_app, jsonify, redirect, render_template, request, session, url_for
from werkzeug.security import check_password_hash, generate_password_hash

from reading_api import MYSQL_CONFIG

auth_bp = Blueprint("auth", __name__)
_DUMMY_HASH = generate_password_hash("not-a-real-password")


def get_account():
    with pymysql.connect(**MYSQL_CONFIG) as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT username, password_hash FROM weread_auth WHERE id = 1")
            return cur.fetchone()


def create_account(username, password_hash):
    try:
        with pymysql.connect(**MYSQL_CONFIG) as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "INSERT INTO weread_auth (id, username, password_hash) VALUES (1, %s, %s)",
                    (username, password_hash),
                )
            conn.commit()
        return True
    except pymysql.IntegrityError:
        return False


def valid_csrf(token):
    expected = session.get("csrf_token", "")
    return bool(expected and token and hmac.compare_digest(expected.encode(), token.encode()))


def csrf_token():
    if not current_app.secret_key:
        return ""
    if "csrf_token" not in session:
        session["csrf_token"] = secrets.token_urlsafe(32)
    return session["csrf_token"]


def start_session(username):
    session.clear()
    session.permanent = True
    session["auth_user"] = username
    csrf_token()


def database_error():
    return render_template(
        "login.html", error="登录数据库暂时不可用，请稍后重试。", unavailable=True,
    ), 503


@auth_bp.get("/login")
def login_page():
    if not current_app.secret_key:
        return render_template(
            "login.html", error="服务未配置会话密钥，请联系管理员。", unavailable=True,
        ), 503
    try:
        initialized = get_account() is not None
    except pymysql.MySQLError:
        return database_error()
    return render_template(
        "login.html",
        initialized=initialized,
        setup_enabled=bool(current_app.config.get("AUTH_SETUP_TOKEN")),
        error=request.args.get("error", ""),
    )


@auth_bp.post("/auth/initialize")
def initialize_account():
    if not valid_csrf(request.form.get("csrf_token")):
        return render_template(
            "login.html", initialized=False, setup_enabled=bool(current_app.config.get("AUTH_SETUP_TOKEN")),
            error="页面已过期，请刷新后重试。",
        ), 400
    try:
        if get_account() is not None:
            return redirect(url_for("auth.login_page"))
    except pymysql.MySQLError:
        return database_error()

    setup_token = current_app.config.get("AUTH_SETUP_TOKEN", "")
    supplied_token = request.form.get("setup_token", "")
    username = request.form.get("username", "").strip()
    password = request.form.get("password", "")
    if not setup_token or not hmac.compare_digest(setup_token.encode(), supplied_token.encode()):
        return render_template("login.html", initialized=False, setup_enabled=bool(setup_token), error="初始化口令无效或尚未配置。"), 403
    if not username or len(username) > 64 or len(password) < 12 or len(password) > 256:
        return render_template(
            "login.html", initialized=False, setup_enabled=True,
            error="用户名需为 1-64 个字符，密码需为 12-256 个字符。",
        ), 400

    try:
        if not create_account(username, generate_password_hash(password)):
            return redirect(url_for("auth.login_page"))
    except pymysql.MySQLError:
        return database_error()

    current_app.config["AUTH_SETUP_TOKEN"] = ""
    start_session(username)
    return redirect("/")


@auth_bp.post("/auth/login")
def login():
    if not valid_csrf(request.form.get("csrf_token")):
        return render_template("login.html", initialized=True, error="页面已过期，请刷新后重试。"), 400
    username = request.form.get("username", "").strip()
    password = request.form.get("password", "")
    if len(password) > 256:
        return render_template("login.html", initialized=True, error="账号或密码不正确。"), 401
    try:
        account = get_account()
    except pymysql.MySQLError:
        return database_error()

    stored_hash = account["password_hash"] if account else _DUMMY_HASH
    password_matches = check_password_hash(stored_hash, password)
    username_matches = account and hmac.compare_digest(account["username"].encode(), username.encode())
    if not account or not username_matches or not password_matches:
        return render_template("login.html", initialized=True, error="账号或密码不正确。"), 401
    start_session(account["username"])
    return redirect("/")


@auth_bp.post("/auth/logout")
def logout():
    if not valid_csrf(request.form.get("csrf_token")):
        return "页面已过期，请刷新后重试。", 400
    session.clear()
    return redirect(url_for("auth.login_page"))


def register_auth(app):
    app.config.update(
        SECRET_KEY=os.environ.get("WEREAD_SESSION_SECRET") or None,
        AUTH_SETUP_TOKEN=os.environ.get("WEREAD_AUTH_SETUP_TOKEN", ""),
        PERMANENT_SESSION_LIFETIME=timedelta(days=5),
        SESSION_REFRESH_EACH_REQUEST=False,
        SESSION_COOKIE_HTTPONLY=True,
        SESSION_COOKIE_SAMESITE="Lax",
        SESSION_COOKIE_SECURE=(
            os.environ.get("WEREAD_SESSION_COOKIE_SECURE", "1").lower()
            not in {"0", "false", "no"}
        ),
    )
    app.register_blueprint(auth_bp)
    app.context_processor(lambda: {"csrf_token": csrf_token})

    @app.before_request
    def require_login_and_csrf():
        if request.endpoint == "static":
            return None
        if not app.secret_key:
            if request.endpoint == "auth.login_page":
                return None
            if request.path.startswith("/api/"):
                return jsonify(error="服务未配置会话密钥"), 503
            return render_template(
                "login.html", unavailable=True, error="服务未配置会话密钥，请联系管理员。",
            ), 503
        if request.endpoint in {"auth.login_page", "auth.login", "auth.initialize_account", "auth.logout"}:
            return None
        if not session.get("auth_user"):
            if request.path.startswith("/api/"):
                return jsonify(error="需要登录"), 401
            return redirect(url_for("auth.login_page"))
        if request.method in {"POST", "PUT", "PATCH", "DELETE"}:
            if not valid_csrf(request.headers.get("X-CSRF-Token")):
                if request.path.startswith("/api/"):
                    return jsonify(error="CSRF token 无效"), 400
                return "页面已过期，请刷新后重试。", 400
        return None