workbuddy-portal/workbuddy_portal/security.py

# -*- coding: utf-8 -*-
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Wang Chuanli

"""鉴权、口令策略、图形验证码、限速、CSRF、安全响应头。

多用户下的三条红线
------------------
1. **Cookie（账号凭证）是按 user_id 分作用域存的**，读取一律走
   `db.get_secret(conn, "cookie", uid)`；`db.get_settings()` 会把凭证置空，
   所以「顺手把配置回传给前端」这类代码不可能把它带出去。
2. **禁用/删除账号立刻失效**：`current_user()` 每个请求回查一次
   users.status 与 users.session_ver，不靠会话过期来兜底。
3. **限速按「来源 IP」和「用户名」双维度计数**，但两者的**强度刻意不同**：
   IP 维度是真锁，用户名维度只是秒级退避。原因见 `user_soft_left` 的注释。

来源 IP 的取法（对外提供服务时最容易出错的一处）
------------------------------------------------
`client_ip()` 是**全站唯一**的取客户端地址入口。默认只信 `remote_addr`：
反向代理若用 `$proxy_add_x_forwarded_for`（追加语义），请求头里第 0 段就是
攻击者自己填的字符串，采信它等于把验证码限速、注册配额、登录锁定三道
IP 防线一起交出去。只有显式设置 `WB_TRUST_PROXY=1`（且你的代理会重写该头）
时才读 X-Forwarded-For，而且**取最右侧**那一段 —— 最右边是离我们最近的
一跳，由我们自己的代理写入，客户端伪造不了。
"""
import functools
import hmac
import ipaddress
import logging
import re
import secrets
import sqlite3
import time

from flask import (current_app, flash, g, jsonify, redirect, render_template,
                   request, session, url_for)

from . import captcha, config, db

log = logging.getLogger("wb.security")

# ---------------- 来源 IP（全站唯一入口） ----------------
def _valid_ip(s):
    try:
        ipaddress.ip_address(s)
        return True
    except ValueError:
        return False


def client_ip():
    """当前请求的客户端地址。

    * `WB_TRUST_PROXY` 未开启（默认）：直接用 `remote_addr`。
      直接暴露公网、或前面挂了「追加型」代理时，XFF 的第一段是攻击者可控的。
    * 已开启：读 X-Forwarded-For 并**取最右侧**合法 IP。
      最右侧是最近一跳（我们自己的代理）写入的，客户端加不进去。
      多级代理（CDN -> nginx）需要按跳数取值，本项目不支持 —— 那样只能靠
      代理侧传 `X-Real-IP` 之类的可信头，不要在这里猜。
    """
    remote = (request.remote_addr or "").strip()
    if not config.TRUST_PROXY:
        return remote
    raw = request.headers.get("X-Forwarded-For", "")
    if not raw:
        return remote
    for part in reversed([p.strip() for p in raw.split(",")]):
        # 去掉 IPv6 的 [..]:port 写法
        cand = part.strip("[]").split("%")[0]
        if cand.count(":") == 1 and cand.rsplit(":", 1)[1].isdigit():
            cand = cand.rsplit(":", 1)[0]         # IPv4:port
        if _valid_ip(cand):
            return cand
    # 头里全是垃圾 -> 退回 remote_addr，而不是把一个伪造值当 IP 用
    log.warning("X-Forwarded-For 里没有合法 IP，已回退 remote_addr：%r", raw[:120])
    return remote


# ---------------- 失败计数（内存即可） ----------------
# 单进程部署（见 README 的部署约束），重启清零可接受；
# 真正的防爆破靠「验证码 + 双维度限速」两道，而不是靠计数持久化。
_fails = {}                  # key -> [count, last_ts]
_tries = {}                  # ip  -> [count, window_started_at]（含成功，只看总量）
_FAILS_MAX_KEYS = 8192       # 上限，防止海量来源把字典撑爆
_FAILS_TTL = 3600            # 超过 1 小时无更新即清理

CAPTCHA_SESSION_PREFIX = "cap_"


def _prune_fails(now=None):
    now = now or time.time()
    dead = [k for k, c in _fails.items() if now - c[1] > _FAILS_TTL]
    for k in dead:
        _fails.pop(k, None)
    if len(_fails) > _FAILS_MAX_KEYS:
        for k, _ in sorted(_fails.items(), key=lambda kv: kv[1][1])[:len(_fails) - _FAILS_MAX_KEYS]:
            _fails.pop(k, None)


def _ip_key(ip):
    return "ip:" + (ip or "")


def _user_key(username):
    return "user:" + (username or "").strip().lower()


def note_fail(key):
    now = time.time()
    _prune_fails(now)
    c = _fails.get(key)
    if c is None or now - c[1] > config.LOGIN_LOCK_MINUTES * 60:
        _fails[key] = [1, now]
        return 1
    c[0] += 1
    c[1] = now
    return c[0]


def is_locked(key):
    c = _fails.get(key)
    if not c or c[0] < config.MAX_LOGIN_FAILS:
        return False
    return time.time() - c[1] <= config.LOGIN_LOCK_MINUTES * 60


def clear_fail(key):
    _fails.pop(key, None)


def lock_left(key):
    c = _fails.get(key)
    if not c:
        return 0
    return max(0, int(config.LOGIN_LOCK_MINUTES * 60 - (time.time() - c[1])))


def fail_count(key):
    c = _fails.get(key)
    return c[0] if c else 0


# ---- IP 维度：真锁（来源地址现在已经不可伪造，锁得住真正的攻击者）----
def ip_lock_left(ip):
    c = _fails.get(_ip_key(ip))
    if not c or c[0] < config.MAX_LOGIN_FAILS:
        return 0
    return max(0, int(config.LOGIN_LOCK_MINUTES * 60 - (time.time() - c[1])))


# ---- 用户名维度：只做秒级退避，**不做长锁** ----
def user_soft_left(username):
    """知道一个用户名就能把它锁死 10 分钟 —— 那本身就是攻击。

    对外提供服务后，管理员用户名是公开信息（导航里就写着），
    如果按用户名施加长锁，任何人只要连打 5 次错误口令，就能让真正的管理员
    十分钟进不去。所以这里改成「递增且有封顶」的秒级等待：
        超过阈值后第 1 次 1s、第 2 次 2s …… 封顶 60s。
    真正的重锁只按来源 IP 施加（`ip_lock_left`）—— 那才是攻击者无法伪造、
    也无法甩锅给别人的东西。命中阈值的同时，攻击者自己的 IP 也在计数，
    所以这种「软」不会让爆破变得可行。
    """
    c = _fails.get(_user_key(username))
    if not c or c[0] < config.USER_SOFT_THRESHOLD:
        return 0
    delay = min(config.USER_SOFT_CAP_SECONDS,
                1 << min(10, c[0] - config.USER_SOFT_THRESHOLD))
    return max(0, int(delay - (time.time() - c[1])))


# ---- 单 IP 登录尝试总量（含成功）：挡住「慢慢撞、不触发失败阈值」----
def note_try(ip):
    now = time.time()
    if len(_tries) > _FAILS_MAX_KEYS:
        _tries.clear()
    cur = _tries.get(ip)
    if cur is None or now - cur[1] > config.LOGIN_ATTEMPTS_WINDOW:
        _tries[ip] = [1, now]
        return 1
    cur[0] += 1
    return cur[0]


def try_window_left(ip):
    cur = _tries.get(ip)
    if not cur or cur[0] < config.LOGIN_ATTEMPTS_PER_IP:
        return 0
    return max(0, int(config.LOGIN_ATTEMPTS_WINDOW - (time.time() - cur[1])))


def auth_locked(ip, username=""):
    """还需等待的秒数（0 = 放行）。"""
    return max(ip_lock_left(ip), try_window_left(ip), user_soft_left(username))


def auth_block_reason(ip, username=""):
    """被挡的原因码：ip / rate / user / 空。用于给出**准确**的提示语。"""
    if ip_lock_left(ip):
        return "ip"
    if try_window_left(ip):
        return "rate"
    if user_soft_left(username):
        return "user"
    return ""


def auth_block_message(reason, seconds):
    if reason == "ip":
        return "该来源登录失败次数过多，请 %d 秒后再试" % seconds
    if reason == "rate":
        return "登录请求过于频繁，请 %d 秒后再试" % seconds
    return "尝试过于频繁，请 %d 秒后再试" % seconds


def note_auth_fail(ip, username=""):
    n1 = note_fail(_ip_key(ip))
    n2 = note_fail(_user_key(username)) if username else 0
    return max(n1, n2)


def clear_auth_fail(ip, username=""):
    """登录成功后清掉两个维度的计数。

    用户名维度必须在成功时清零：否则「攻击者打了几次 + 主人自己登一次」
    之后，主人仍会被自己之前那几次的退避拖住。
    """
    clear_fail(_ip_key(ip))
    if username:
        clear_fail(_user_key(username))


# ---------------- 通用动作限速（重操作保护） ----------------
# 采集 / 导出 / 整库整理这类动作的代价远高于普通页面请求：
# 一次 /api/collect 会让服务端对云端发起成百上千次请求，并独占一个线程；
# 一次导出会把整张表扫一遍。对外提供服务时必须有刹车，
# 否则**一个注册账号**就能把实例的线程与带宽吃干净。
_actions = {}                 # "name:uid" -> next_allowed_ts


def action_allowed(key, min_interval):
    """返回 (是否允许, 还需等待秒数)。允许时会把下次可执行时间推后。"""
    now = time.time()
    if len(_actions) > _FAILS_MAX_KEYS:
        _actions.clear()
    nxt = _actions.get(key) or 0
    if now < nxt:
        return False, int(nxt - now) + 1
    _actions[key] = now + max(0, int(min_interval))
    return True, 0


def action_wait_left(key):
    return max(0, int((_actions.get(key) or 0) - time.time()))


# ---------------- 口令 / 用户名策略 ----------------
_USERNAME_RE = re.compile(config.USERNAME_RE)


def hash_password(p):
    from werkzeug.security import generate_password_hash
    return generate_password_hash(p, method="pbkdf2:sha256:200000")


def verify_password(hashed, p):
    from werkzeug.security import check_password_hash
    try:
        return check_password_hash(hashed, p)
    except (ValueError, TypeError):
        return False


def username_problem(name):
    """校验用户名。开放注册后这是第一个入口，必须收紧。"""
    name = (name or "").strip()
    if not name:
        return "用户名必填"
    if not _USERNAME_RE.match(name):
        return "用户名需 3~32 位，以字母或数字开头，只能用字母、数字、下划线、点、连字符"
    if name.lower() in ("admin", "administrator", "root", "system", "guest", "null"):
        return "该用户名为系统保留字，请换一个"
    return None


def password_problem(new, new2=None, username=None):
    """口令强度：8 位以上，且至少包含两类字符。

    比原来的「只要 6 位」严格——因为现在任何人都能自助注册，
    弱口令直接决定了整个实例的抗爆破能力。
    """
    new = new or ""
    if len(new) < config.PASSWORD_MIN:
        return "密码至少 %d 位" % config.PASSWORD_MIN
    if len(new) > config.PASSWORD_MAX:
        return "密码过长（上限 %d 位）" % config.PASSWORD_MAX
    classes = sum(bool(re.search(p, new)) for p in
                  (r"[a-z]", r"[A-Z]", r"[0-9]", r"[^A-Za-z0-9]"))
    if classes < 2:
        return "密码需包含大写字母、小写字母、数字、符号中的至少两类"
    # 黑名单只挡「字典头几页」，命中即拒。只在设置/修改口令时校验，
    # 登录路径不校验 —— 所以不会把用老口令的人挡在门外。
    if new.lower() in config.WEAK_PASSWORDS:
        return "这个密码在常见弱口令字典里，请换一个"
    if new2 is not None and new2 != new:
        return "两次输入的新密码不一致"
    if username and new.lower() == str(username).lower():
        return "密码不能与用户名相同"
    return None


# 兼容旧名（原来的 api.py 内部函数）
_check_password = password_problem


# ---------------- 登录 ----------------
def login_ok(conn, username, password):
    """校验口令。返回 (user_row, error_message)。

    停用账号与口令错误返回**同一句话**，避免探测哪些用户名存在
    （不过自助注册本身就暴露了用户名唯一性，这里只是不打额外的广告）。
    """
    username = (username or "").strip()
    row = conn.execute("SELECT * FROM users WHERE username=?", (username,)).fetchone()
    if row is None or not verify_password(row["password_hash"], password):
        return None, "用户名或密码不正确"
    if (row["status"] or "active") != "active":
        return None, "该账号已被停用，请联系管理员"
    conn.execute("UPDATE users SET last_login_at=?, login_count=login_count+1 WHERE id=?",
                 (db.now_str(), row["id"]))
    return row, None


# ---------------- 会话 ----------------
def current_user():
    """当前登录用户（dict）或 None。

    每个请求回查一次 users 表，两道校验：
      * `status` —— 账号被停用/删除后**立刻**失效，而不是等 12 小时会话过期
      * `session_ver` —— 改密码 / 管理员重置 / 停用后，签发时的那一版会话
        立即作废。少了它，「怀疑会话泄漏了所以改密码」就是个假的安心动作：
        旧会话照样有效到 12 小时之后。
    结果缓存在 flask.g 里，一次请求只查一次。
    """
    if "wb_user" in g:
        return g.wb_user
    uid = session.get("uid")
    user = None
    if uid:
        try:
            row = db.get_db().execute(
                "SELECT id,username,display_name,is_admin,status,session_ver"
                " FROM users WHERE id=?", (uid,)).fetchone()
        except sqlite3.OperationalError as e:
            # users 表结构与代码不一致（典型场景：升级到新版本后没跑 init_db
            # 就先把 Web 起起来了，老库还没有 session_ver 这一列）。
            # 这里必须**大声**报错。若和下面「无请求上下文」一起被静默吞掉，
            # 症状会变成「全站所有人被踢下线、日志里什么都没有」——
            # 界面上只看到「登录成功又立刻跳回登录页」，极难归因。
            log.error("会话校验失败：users 表结构与代码不一致（%s）；"
                      "请执行 `python manage.py init` 完成迁移", e)
            row = None
        except Exception:                       # noqa: BLE001  (无请求上下文等)
            row = None
        if row is None or (row["status"] or "active") != "active":
            session.clear()
        elif session.get("sv", 0) != db.session_ver_of(row):
            # 老会话没有 sv 字段时按 0 处理，这样升级本身不会把所有人踢下线
            session.clear()
        else:
            user = {"id": row["id"], "username": row["username"],
                    "display_name": row["display_name"] or row["username"],
                    "is_admin": bool(row["is_admin"])}
            # 个人信息（显示名）改过之后立即生效，不必重新登录
            session["dname"] = user["display_name"]
            session["adm"] = 1 if user["is_admin"] else 0
    g.wb_user = user
    return user


def is_admin():
    u = current_user()
    return bool(u and u.get("is_admin"))


def login_session(user):
    """建立登录会话。

    `session.clear()` 是必须的：既清掉前一次的残留，
    也顺带换掉 CSRF token 与验证码 id —— 这正是防「会话固定」的做法。

    `sv` 记下签发时的 users.session_ver：之后一旦账号改密 / 被重置 / 被停用，
    这一版会话会在下一个请求就被判为过期。
    """
    session.clear()
    session["uid"] = user["id"]
    session["uname"] = user["username"]
    session["dname"] = user["display_name"] or user["username"]
    session["adm"] = 1 if user["is_admin"] else 0
    session["sv"] = db.session_ver_of(user)
    session["login_at"] = db.now_str()
    session.permanent = True


def logout_session():
    session.clear()


def wants_json():
    return (request.path.startswith("/api/")
            or request.accept_mimetypes.best == "application/json")


def login_required(fn):
    @functools.wraps(fn)
    def wrapper(*a, **kw):
        if current_user() is None:
            if wants_json():
                return jsonify({"ok": False, "error": "unauthorized",
                                "message": "登录已失效，请重新登录"}), 401
            return redirect(url_for("views.login", next=request.full_path))
        return fn(*a, **kw)
    return wrapper


def admin_required(fn):
    """管理员专属操作（用户管理等）。非管理员返回 403。"""
    @functools.wraps(fn)
    @login_required
    def wrapper(*a, **kw):
        if not is_admin():
            if wants_json():
                return jsonify({"ok": False, "error": "forbidden",
                                "message": "只有管理员可以执行该操作"}), 403
            return render_template("error.html", code=403, message="只有管理员可以执行该操作"), 403
        return fn(*a, **kw)
    return wrapper


# ---------------- 跳转目标白名单（防开放重定向） ----------------
def safe_next(target, fallback="/"):
    """只允许站内相对路径。

    `//evil.com`、`/\\evil.com`、`https://evil.com` 都必须拒绝：
    `//` 开头是协议相对 URL，浏览器会把 `//evil.com` 当成外站跳转。
    """
    if not target:
        return fallback
    t = str(target).strip()
    if not t.startswith("/"):
        return fallback
    if t.startswith("//") or t.startswith("/\\") or "\\" in t:
        return fallback
    if t.lstrip("/").startswith("//"):
        return fallback
    if "\r" in t or "\n" in t:
        return fallback
    return t


# ---------------- 图形验证码 ----------------
def captcha_required(conn, ip, username=""):
    """按 captcha_policy 决定本次是否需要验证码。"""
    policy = (db.get_setting(conn, "captcha_policy", "always") or "always").strip().lower()
    if policy == "off":
        return False
    if policy == "adaptive":
        # 「自适应」= 这个来源出过问题才要求，日常登录不打扰
        return (fail_count(_ip_key(ip)) >= 2
                or (username and fail_count(_user_key(username)) >= 2))
    return True                             # always（默认）


def issue_captcha(conn, purpose):
    """新建验证码并把 id 记进会话，返回 PNG 字节。答案绝不离开服务端。"""
    try:
        length = int(db.get_setting(conn, "captcha_length", 4) or 4)
    except (TypeError, ValueError):
        length = 4
    length = max(4, min(6, length))
    cid, code = captcha.create(conn, purpose, length=length)
    session[CAPTCHA_SESSION_PREFIX + purpose] = cid
    return captcha.render(code, width=150 if length <= 4 else 150 + (length - 4) * 32)


def consume_captcha(conn, purpose, answer):
    """校验并作废本次验证码。会话里的 id 一并丢掉，逼迫下次换一张新图。"""
    cid = session.pop(CAPTCHA_SESSION_PREFIX + purpose, None)
    return captcha.verify(conn, cid, (answer or "").strip().upper(), purpose)


# 验证码出图限速：出图本身要做点阵渲染 + zlib，不设限就是一条廉价的
# CPU/带宽放大路径（有人拿它当免费的图片生成器刷）。
_cap_fetch = {}                  # ip -> [count, window_started_at]
_CAP_FETCH_MAX = 40              # 每窗口最多出图张数
_CAP_FETCH_WINDOW = 60           # 窗口长度（秒）


def captcha_fetch_allowed(ip):
    now = time.time()
    cur = _cap_fetch.get(ip)
    if cur is None or now - cur[1] > _CAP_FETCH_WINDOW:
        if len(_cap_fetch) > _FAILS_MAX_KEYS:
            _cap_fetch.clear()
        _cap_fetch[ip] = [1, now]
        return True
    cur[0] += 1
    return cur[0] <= _CAP_FETCH_MAX


def audit_login_fail(conn, username, detail, ip):
    """登录失败审计。

    `user_id` 留 0：此时还不能确定是谁（可能是有人在撞别人的账号），
    但 `actor` 记下被尝试的用户名，便于事后按人名检索。
    """
    db.audit(conn, "login_failed", username or "-", detail, ip, 0)


# ---------------- 注册开关与配额 ----------------
def register_allowed(conn):
    return db.get_bool(conn, "allow_register", True)


def register_quota(conn, ip):
    """同一 IP 当天的注册配额。返回 (是否允许, 已注册数, 上限)。"""
    limit = db.get_int(conn, "register_max_per_ip", 3)
    today = db.now_str()[:10]
    n = conn.execute("SELECT COUNT(*) FROM users WHERE register_ip=?"
                     " AND substr(COALESCE(created_at,''),1,10)=?", (ip, today)).fetchone()[0]
    return n < limit, n, limit


# ---------------- CSRF ----------------
def csrf_token():
    t = session.get("_csrf")
    if not t:
        t = session["_csrf"] = secrets.token_urlsafe(24)
    return t


def check_csrf():
    """对所有 POST/PUT/DELETE 生效，失败直接 400。"""
    if request.method in ("GET", "HEAD", "OPTIONS"):
        return None
    sent = request.form.get("_csrf") or request.headers.get("X-CSRF-Token") or ""
    if not sent or not hmac.compare_digest(sent, session.get("_csrf", "")):
        if wants_json():
            return jsonify({"ok": False, "error": "csrf",
                            "message": "CSRF 校验失败，请刷新页面"}), 400
        return "CSRF 校验失败，请刷新页面后重试", 400
    return None


# ---------------- 安全响应头 ----------------
# 这些头是「纵深防御」：本项目的输出都过了 Jinja 自动转义 + app.js 手动转义，
# 但多一层 nosniff / frame-ancestors 能让「某处漏转义」不至于直接变成可利用的 XSS。
CSP = ("default-src 'self'; "
       "img-src 'self' data:; "
       "style-src 'self' 'unsafe-inline'; "
       "script-src 'self' 'unsafe-inline'; "
       "connect-src 'self'; "
       "font-src 'self' data:; "
       "object-src 'none'; "
       "base-uri 'self'; "
       "form-action 'self'; "
       "frame-ancestors 'none'")


def apply_security_headers(resp):
    resp.headers.setdefault("X-Content-Type-Options", "nosniff")
    resp.headers.setdefault("X-Frame-Options", "DENY")
    # 不让站内 URL（可能含 next=、run= 等参数）随外链 referer 泄漏出去
    resp.headers.setdefault("Referrer-Policy", "same-origin")
    resp.headers.setdefault("Content-Security-Policy", CSP)
    resp.headers.setdefault("Cross-Origin-Opener-Policy", "same-origin")
    # HSTS 只在「确认这个部署跑在 HTTPS 上」时才发：在纯 HTTP 部署上发它，
    # 浏览器会把该域名的 http 访问强行升级，表现成「打开就白屏」。
    # 判据是管理员显式打开了 COOKIE_SECURE 或 FORCE_HTTPS。
    if config.COOKIE_SECURE or config.FORCE_HTTPS:
        resp.headers.setdefault("Strict-Transport-Security",
                                "max-age=31536000; includeSubDomains")
    if request.path.startswith("/api/") or request.path.startswith("/captcha"):
        resp.headers.setdefault("Cache-Control", "no-store")
    return resp


def needs_https_redirect():
    """当前请求是否该被跳到 https（仅在显式开启 WB_FORCE_HTTPS 时才判断）。"""
    if not config.FORCE_HTTPS or request.is_secure:
        return False
    # 反代终止 TLS 时，Flask 看到的是 http；靠 X-Forwarded-Proto 还原真实协议。
    # 这个头只在「你已经决定信任代理」的前提下才有意义，所以与 TRUST_PROXY 绑定。
    if config.TRUST_PROXY and (request.headers.get("X-Forwarded-Proto") or "").lower() == "https":
        return False
    if request.method not in ("GET", "HEAD"):
        return False                       # 不重定向 POST：会丢请求体，行为难以预期
    return True


def _access_log(resp, started):
    if not config.ACCESS_LOG:
        return resp
    path = request.path
    if path.startswith("/static/") or path == "/captcha.png":
        return resp                        # 静态资源与验证码出图会把日志刷满
    log.info("%s %s -> %s  %dms  ip=%s", request.method, path, resp.status_code,
             int((time.time() - started) * 1000), client_ip())
    return resp


def init_app(app):
    app.jinja_env.globals["csrf_token"] = csrf_token
    app.jinja_env.globals["current_user"] = current_user

    @app.before_request
    def _guard():
        g.wb_t0 = time.time()
        if needs_https_redirect():
            url = request.url.replace("http://", "https://", 1)
            return redirect(url, code=301)
        return check_csrf()

    @app.after_request
    def _headers(resp):
        _access_log(resp, getattr(g, "wb_t0", time.time()))
        return apply_security_headers(resp)