NeteaseUrl/auth_system/users.py

# -*- coding: utf-8 -*-
"""用户增删改查、登录校验与状态。

（自 auth_system.py 拆出，行为与原实现一致；DB 访问已迁移为原生异步。）
"""

import asyncio
import json
import sqlite3
import time
from dataclasses import dataclass
from typing import Any

from .db import (
    _q,
    _x,
    write_transaction,
)
from .security import (
    _hash_new,
    _normalize_email,
    _verify_password,
)


@dataclass(frozen=True)
class RegistrationGuard:
    """注册风控上下文：同 IP 频率闸 + 设备赠送闸。

    这些检查必须与建号 INSERT 在**同一写事务内**完成——锁外预检存在
    TOCTOU（并发同设备双注册各得一次赠送、并发同 IP 挤过限额），所以它们不是
    「可选的附加参数」，而是事务内必须可见的输入。管理建号无风控需求，不传即可。
    """

    device_id: str = ''
    ip: str = ''
    ua: str = ''
    ip_limit: int = 0
    ip_window: int = 86400
    device_once: bool = False


async def create_user(username: str, password: str, email: str = '', credits: int = 0,
                      role: str = 'user', *,
                      guard: RegistrationGuard | None = None) -> sqlite3.Row | None:
    # 余额只接受 -1（不限）或 >= 0；非法负值一律按 0 处理，禁止出现 -2 这类"假不限"
    try:
        credits = int(credits)
    except (TypeError, ValueError):
        credits = 0
    if credits < -1:
        credits = 0
    # 注册入口传 guard=RegistrationGuard(ip_limit=3, device_once=True)——同 IP 24h
    # 注册数检查在写事务内与 INSERT 原子完成，堵住「锁外检查 + 并发插入」
    # 的 TOCTOU 刷号；管理建号不传 guard（即不限制）。超限抛 ValueError 由调用方转 400/429。
    guard = guard or RegistrationGuard()
    now = int(time.time())
    # PBKDF2 10 万次迭代是 CPU 密集（~50ms）：在事务外先算好，避免持写锁时
    # 阻塞事件循环（正当的 CPU 卸载，非同步门面）。
    phash = await asyncio.to_thread(_hash_new, password)
    async with write_transaction() as conn:
        if guard.ip_limit > 0 and guard.ip:
            cur = await conn.execute(
                'SELECT COUNT(*) AS n FROM users WHERE register_ip = ? AND created_at > ?',
                (guard.ip, now - guard.ip_window))
            row = await cur.fetchone()
            n = int(row['n']) if row else 0
            if n >= guard.ip_limit:
                raise ValueError('当前网络注册过于频繁，请 24 小时后再试或联系管理员')
        if guard.device_once and guard.device_id:
            # 设备赠送闸并入同一事务——锁外 device_used 预检有 TOCTOU，
            # 并发同设备双注册可各得一次赠送；锁内复查命中则本次建号不赠送（账号照常创建）
            cur = await conn.execute(
                "SELECT COUNT(*) AS n FROM users WHERE trial_device = ? AND role != 'admin'",
                (guard.device_id,))
            row = await cur.fetchone()
            if row and int(row['n']) > 0:
                credits = 0
        await conn.execute(
            'INSERT INTO users (username, password_hash, salt, email, role, duration_seconds, '
            'trial_device, register_ip, register_ua, last_active, created_at) '
            'VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)',
            (username, phash, '', email, role,
             -1 if role == 'admin' else int(credits),
             guard.device_id, guard.ip, guard.ua, now, now))
    return await get_user_by_name(username)


async def get_user_by_name(username: str) -> sqlite3.Row | None:
    rows = await _q('SELECT * FROM users WHERE username = ?', (username,))
    return rows[0] if rows else None


async def get_user_by_email(email: str) -> sqlite3.Row | None:
    """按邮箱查用户（注册/建号邮箱查重的友好提示；库侧另有唯一索引兜底并发）"""
    email = _normalize_email(email)
    if not email:
        return None
    rows = await _q('SELECT * FROM users WHERE email = ? LIMIT 1', (email,))
    return rows[0] if rows else None


async def register_ip_exceeded(ip: str, limit: int = 3, window: int = 86400) -> bool:
    """同 IP 24h 内注册数超限（默认 3 次）——防 IP 轮换刷免费赠送次数；
    设备指纹可被客户端伪造，IP 维度是第二道闸。家庭 NAT 场景超限可由管理员建号兜底。"""
    if not ip:
        return False
    rows = await _q('SELECT COUNT(*) AS n FROM users WHERE register_ip = ? AND created_at > ?',
                    (ip, int(time.time() - window)))
    return bool(rows and rows[0]['n'] >= limit)


async def get_user_by_id(uid: int) -> sqlite3.Row | None:
    rows = await _q('SELECT * FROM users WHERE id = ?', (uid,))
    return rows[0] if rows else None


async def verify_login(username: str, password: str) -> sqlite3.Row | None:
    """校验用户名密码。密码正确即返回用户行（含被禁用/封禁账号——
    状态语义由登录路由判断：禁用显示"已被禁用"、封禁显示封禁原因。
    原实现把禁用账号提前过滤成 None，路由的"账号已被禁用"分支永不可达）"""
    user = await get_user_by_name(username)
    if user:
        # PBKDF2 校验是 CPU 密集（~50ms），卸载到线程避免阻塞事件循环（正当 CPU 卸载）
        ok, needs_upgrade = await asyncio.to_thread(
            _verify_password, user['password_hash'], password, user['salt'])
        if ok:
            if needs_upgrade:
                # 存量 SHA-256 哈希在登录成功时透明升级为 PBKDF2
                new_hash = await asyncio.to_thread(_hash_new, password)
                await _x('UPDATE users SET password_hash = ?, salt = \'\' WHERE id = ?',
                         (new_hash, user['id']))
            return user
    return None


async def list_users() -> list[dict[str, Any]]:
    rows = await _q('SELECT * FROM users ORDER BY id DESC')
    return [_user_public(r, full=True) for r in rows]


async def delete_user(uid: int) -> bool:
    """删除用户（卡密 used_by 为用户名字符串无外键——删除时同步
    标注"（已注销#id）"，管理端卡密列表不再显示悬空的已删用户名）"""
    u = await get_user_by_id(uid)
    if not u or u['role'] == 'admin':
        return False
    async with write_transaction() as conn:
        name = u['username']
        await conn.execute('DELETE FROM users WHERE id = ? AND role != ?', (uid, 'admin'))
        await conn.execute("UPDATE cards SET used_by = used_by || ? WHERE used_by = ?",
                           (f'（已注销#{uid}）', name))
    return True


async def reset_password(uid: int, new_password: str) -> bool:
    """重置密码（同时递增 login_seq，让所有已签发会话立即失效）"""
    new_hash = await asyncio.to_thread(_hash_new, new_password)
    await _x('UPDATE users SET password_hash = ?, salt = \'\', '
             'login_seq = COALESCE(login_seq, 0) + 1 WHERE id = ?',
             (new_hash, uid))
    return True


async def set_user_status(uid: int, status: int) -> bool:
    """启用 / 禁用账号（1 启用，0 禁用）；管理员账号不可禁用"""
    if status not in (0, 1):
        return False
    row = await _q('SELECT role FROM users WHERE id = ?', (uid,))
    if row and row[0]['role'] == 'admin':
        return False
    return bool(await _x('UPDATE users SET status = ? WHERE id = ?', (status, uid)))


async def ban_user(uid: int, reason: str = '') -> bool:
    """封禁账号（管理员不可封）"""
    row = await _q('SELECT role FROM users WHERE id = ?', (uid,))
    if not row or row[0]['role'] == 'admin':
        return False
    return bool(await _x('UPDATE users SET status = 2, ban_reason = ? WHERE id = ?',
                         (str(reason or '')[:200], uid)))


async def unban_user(uid: int) -> bool:
    """解封账号，恢复正常（清空触发封禁的省份历史，否则下次登录立即再封）"""
    return bool(await _x('UPDATE users SET status = 1, ban_reason = \'\', recent_provinces = \'[]\' WHERE id = ?', (uid,)))


async def set_user_credits(uid: int, count: int) -> None:
    """设置用户剩余次数（-1 表示不限）"""
    count = -1 if count < 0 else int(count)
    await _x('UPDATE users SET duration_seconds = ? WHERE id = ?', (count, uid))


async def set_user_password(uid: int, password: str) -> None:
    # 新哈希 + 递增 login_seq（所有旧会话立即失效）
    new_hash = await asyncio.to_thread(_hash_new, password)
    await _x('UPDATE users SET password_hash = ?, salt = \'\', '
             'login_seq = COALESCE(login_seq, 0) + 1 WHERE id = ?',
             (new_hash, uid))


async def set_user_netease_cookie(username: str, cookie: str) -> None:
    await _x('UPDATE users SET netease_cookie = ? WHERE username = ?', (cookie, username))


def _user_public(r: sqlite3.Row, full: bool = False) -> dict[str, Any]:
    d = {
        'id': r['id'],
        'username': r['username'],
        'email': r['email'],
        'role': r['role'],
        'duration_seconds': r['duration_seconds'],
        'created_at': r['created_at'],
        'last_login': r['last_login'],
        'ann_seen': r['ann_seen'],
    }
    if full:
        d.update({
            'status': r['status'] if 'status' in r.keys() else 1,
            'trial_device': r['trial_device'],
            'register_ip': r['register_ip'],
            'register_ua': r['register_ua'],
            # 列表不再携带明文网易云 Cookie（前端只用布尔值；原文按需走
            # /admin/site/user/cookie 单查接口，避免一次列表请求扩散全部用户凭据）
            'has_netease_cookie': bool(r['netease_cookie']),
            'ban_reason': r['ban_reason'] if 'ban_reason' in r.keys() else '',
            'last_province': _last_province(r),
        })
    return d


def _last_province(r: sqlite3.Row) -> str:
    try:
        history = json.loads(r['recent_provinces'] if 'recent_provinces' in r.keys() and r['recent_provinces'] else '[]')
        return history[-1].get('p', '') if history else ''
    except Exception:
        return ''


# ============ 免费体验（设备识别） ============


async def device_used(device_id: str) -> bool:
    if not device_id:
        return False
    rows = await _q("SELECT id FROM users WHERE trial_device = ? AND role != 'admin' LIMIT 1", (device_id,))
    return bool(rows)