This repository has been archived on 2026-10-02. You can view files and clone it. You cannot open issues or pull requests or push a commit.
Files
NeteaseUrl/auth_system/db.py
T
snowy-yang 37f8d77405 refactor: init_db 四段拆解(142 行 / 圈 15 / 嵌套 4 → 14 行 / 圈 2 / 嵌套 1)
为什么:init_db 是站点库全部结构演进的唯一入口,但 142 行单函数把
「建表 / 列迁移 / 邮箱唯一索引 / 邮箱归一化 / 管理员引导」五件事
揉在一个 with-try 里,任何一段改动都要在 140 行上下文里核对边界。

拆出五个成员(全部接收 conn,调用顺序 = 原执行顺序,单连接单 commit 不变):
- _SCHEMA_SQL 常量 + _create_schema(建表 + 种子行,纯数据上移)
- _apply_column_migrations(ALTER 循环;非 duplicate column 的
  OperationalError 仍然抛出,禁止静默带病初始化)
- _ensure_email_unique_index(重复邮箱仅告警不阻断,原文案保留)
- _migrate_user_emails(撞唯一索引行保留原值,原文案保留)
- _ensure_initial_admin(按 role 判断而非用户名,防改名后误重建)

先写测试后动代码:新增 tests/test_auth_init_db.py 7 例,覆盖此前
零测试的路径——旧库列迁移、二次初始化幂等、改名管理员不重建、
存量重复邮箱告警不阻断、+tag 邮箱归一化迁移、归一化撞行保留原值。
7 例在未改动实现上一次跑绿,重构后即等价性验证。
门禁:ruff 全过 · pyright 0/0 · pytest 319 passed(312→319)。
2026-09-27 09:36:31 +08:00

282 lines
11 KiB
Python

# -*- coding: utf-8 -*-
"""路径常量、连接与建表。
(自 auth_system.py 拆出,行为与原实现一致。)
"""
import os
import secrets
import sqlite3
import sys
import threading
import time
from pathlib import Path
from .security import (
_hash_new,
_normalize_email,
)
# ============ 路径与常量 ============
# 必须与 core.db.DATA_DIR 用**同一条**表达式:两处指向同一个 data/site.db。
# 此前这里写死 Path('data') 而 core 侧认 SITE_DATA_DIR——一旦运维设置该变量,
# 站点库与设置/公告库会落到两个不同文件(脑裂),且不会报错。
DATA_DIR = Path(os.environ.get('SITE_DATA_DIR') or 'data')
DB_PATH = DATA_DIR / 'site.db'
BACKUP_DIR = DATA_DIR / 'backups'
SECRET_KEY_PATH = DATA_DIR / 'secret.key'
DEFAULT_ADMIN_USER = 'admin'
SESSION_KEY = 'site_uid'
DEFAULT_REGISTER_CREDITS = 3 # 注册直接赠送 3 次下载次数(无需领取;TRIAL_CREDITS 死常量已删)
_DB_LOCK = threading.RLock()
# ============ 数据库 ============
def get_conn() -> sqlite3.Connection:
DATA_DIR.mkdir(parents=True, exist_ok=True)
conn = sqlite3.connect(str(DB_PATH), timeout=15)
conn.row_factory = sqlite3.Row
conn.execute('PRAGMA journal_mode=WAL')
return conn
def _q(sql: str, params: tuple = ()) -> list[sqlite3.Row]:
"""单语句读:**不持锁**。
WAL 模式下每个连接读到一致性快照,读与写互不阻塞;旧实现对读也加 _DB_LOCK,
等于把全部读串行化并被写挡住。此处仅放开单语句读——写事务(_x 与各
`with _DB_LOCK` 块)仍全程串行,故不引入写-写竞争。
"""
conn = get_conn()
try:
return conn.execute(sql, params).fetchall()
finally:
conn.close()
def _x(sql: str, params: tuple = ()) -> int:
"""执行写操作,返回 lastrowid / rowcount"""
with _DB_LOCK:
conn = get_conn()
try:
cur = conn.execute(sql, params)
conn.commit()
return cur.lastrowid if cur.lastrowid else cur.rowcount
finally:
conn.close()
_SCHEMA_SQL = '''
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT UNIQUE NOT NULL,
password_hash TEXT NOT NULL,
salt TEXT NOT NULL,
email TEXT DEFAULT '',
role TEXT DEFAULT 'user',
duration_seconds INTEGER DEFAULT 0,
trial_device TEXT DEFAULT '',
register_ip TEXT DEFAULT '',
register_ua TEXT DEFAULT '',
netease_cookie TEXT DEFAULT '',
ann_seen TEXT DEFAULT '',
last_active INTEGER DEFAULT 0,
created_at INTEGER DEFAULT 0,
last_login INTEGER DEFAULT 0
);
CREATE TABLE IF NOT EXISTS email_codes (
id INTEGER PRIMARY KEY AUTOINCREMENT,
email TEXT NOT NULL,
code TEXT NOT NULL,
expires_at INTEGER NOT NULL,
created_at INTEGER DEFAULT 0,
purpose TEXT DEFAULT 'register'
);
CREATE TABLE IF NOT EXISTS cards (
id INTEGER PRIMARY KEY AUTOINCREMENT,
code TEXT UNIQUE NOT NULL,
minutes INTEGER NOT NULL,
created_by TEXT DEFAULT '',
created_at INTEGER DEFAULT 0,
used_by TEXT DEFAULT '',
used_at INTEGER DEFAULT 0
);
CREATE TABLE IF NOT EXISTS logs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
type TEXT DEFAULT 'success',
username TEXT DEFAULT '',
ip TEXT DEFAULT '',
action TEXT DEFAULT '',
detail TEXT DEFAULT '',
time TEXT DEFAULT ''
);
'''
def _create_schema(conn: sqlite3.Connection) -> None:
"""建表 + 单行配置种子行(幂等:全部 IF NOT EXISTS / INSERT OR IGNORE)。"""
conn.executescript(_SCHEMA_SQL)
# 单行配置表
conn.execute('''CREATE TABLE IF NOT EXISTS smtp_config (
id INTEGER PRIMARY KEY CHECK (id = 1),
host TEXT DEFAULT '', port INTEGER DEFAULT 587,
sender TEXT DEFAULT '', password TEXT DEFAULT '',
updated_at INTEGER DEFAULT 0
)''')
conn.execute('INSERT OR IGNORE INTO smtp_config (id, port) VALUES (1, 587)')
# 键值设置表
conn.execute('''CREATE TABLE IF NOT EXISTS settings (
key TEXT PRIMARY KEY, value TEXT DEFAULT ''
)''')
conn.execute('''CREATE TABLE IF NOT EXISTS announcements (
id INTEGER PRIMARY KEY CHECK (id = 1),
title TEXT DEFAULT '', content TEXT DEFAULT '',
popup_mode TEXT DEFAULT 'off',
updated_at INTEGER DEFAULT 0
)''')
conn.execute('INSERT OR IGNORE INTO announcements (id) VALUES (1)')
# 计费账本(任务实扣/交付/退款/终态的持久化唯一记录,结算幂等)
conn.execute('''CREATE TABLE IF NOT EXISTS billing_ledger (
task_id TEXT PRIMARY KEY,
user_id INTEGER NOT NULL,
reserved INTEGER DEFAULT 0,
delivered INTEGER DEFAULT 0,
refunded INTEGER DEFAULT 0,
status TEXT DEFAULT 'open',
created_at INTEGER DEFAULT 0,
settled_at INTEGER DEFAULT 0
)''')
def _apply_column_migrations(conn: sqlite3.Connection) -> None:
"""轻量迁移:新增列。仅"列已存在"可忽略(其余迁移失败必须暴露,不能吞掉)"""
_already: set[str] = set()
for trow in conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall():
for crow in conn.execute(f"PRAGMA table_info({trow['name']})").fetchall():
_already.add(f"{trow['name']}.{crow['name']}")
for alter_sql, tbl_col in (
('ALTER TABLE users ADD COLUMN status INTEGER DEFAULT 1', 'users.status'),
('ALTER TABLE users ADD COLUMN login_seq INTEGER DEFAULT 0', 'users.login_seq'),
('ALTER TABLE users ADD COLUMN recent_provinces TEXT DEFAULT \'[]\'', 'users.recent_provinces'),
('ALTER TABLE users ADD COLUMN ban_reason TEXT DEFAULT \'\'', 'users.ban_reason'),
('ALTER TABLE cards ADD COLUMN note TEXT DEFAULT \'\'', 'cards.note'),
('ALTER TABLE cards ADD COLUMN expire_at INTEGER DEFAULT 0', 'cards.expire_at'),
('ALTER TABLE cards ADD COLUMN disabled INTEGER DEFAULT 0', 'cards.disabled'),
('ALTER TABLE email_codes ADD COLUMN purpose TEXT DEFAULT \'register\'', 'email_codes.purpose'),
):
if tbl_col in _already:
continue # 幂等:列已存在
try:
conn.execute(alter_sql)
except sqlite3.OperationalError as e:
if 'duplicate column' not in str(e).lower():
# 磁盘/锁/结构等真实故障:抛出,禁止静默带病初始化
raise
def _ensure_email_unique_index(conn: sqlite3.Connection) -> None:
"""非空邮箱唯一(部分索引,不约束历史空串)。
存量库若已有重复非空邮箱,索引创建失败仅告警不阻断启动
(查重逻辑在注册/建号路由另有兜底)。
"""
try:
conn.execute('CREATE UNIQUE INDEX IF NOT EXISTS idx_users_email_nonempty '
'ON users(email) WHERE email != \'\'')
except sqlite3.IntegrityError:
print('[警告] users.email 存在重复的非空邮箱,唯一索引未创建(请人工清理后重启)')
def _migrate_user_emails(conn: sqlite3.Connection) -> None:
"""users.email 口径统一为归一化(lower + 剥 Gmail +tag)——与
get_user_by_email/找回/邮箱登录的查询口径对齐。逐行迁移,归一化后与他
人撞唯一索引的行保留原值跳过(与"存量重复仅告警"同口径)。幂等:已
归一化的行零写入。
"""
for urow in conn.execute("SELECT id, email FROM users WHERE email != ''").fetchall():
fixed = _normalize_email(urow['email'])
if not fixed or fixed == urow['email']:
continue
try:
conn.execute('UPDATE users SET email = ? WHERE id = ?', (fixed, urow['id']))
except sqlite3.IntegrityError:
print(f"[警告] 用户 {urow['id']} 的邮箱归一化后与他人重复,保留原值")
def _ensure_initial_admin(conn: sqlite3.Connection) -> None:
"""初始管理员(新装不再使用公开默认口令,随机生成并写入日志一次)。
引导条件改为"不存在任何管理员角色"——按用户名查会在
管理员改名(如 admin→qfbsh)后误判缺失,重启即重建出一个新的 admin 账号。
"""
row = conn.execute("SELECT id FROM users WHERE role = 'admin' LIMIT 1").fetchone()
if row:
return
initial_pass = secrets.token_urlsafe(12)
conn.execute(
'INSERT INTO users (username, password_hash, salt, role, duration_seconds, created_at) '
'VALUES (?, ?, ?, ?, -1, ?)',
(DEFAULT_ADMIN_USER, _hash_new(initial_pass), '',
'admin', int(time.time())))
# ——非 ASCII 输出须防 GBK 重定向下 UnicodeEncodeError
try:
print(f'[初始化] 已创建管理员账号 admin,初始密码:{initial_pass}(请立即登录并修改,本提示不再显示)')
except UnicodeEncodeError:
print('[init] admin account created, initial password printed above failed to encode')
def init_db() -> None:
"""初始化表结构与默认数据"""
DATA_DIR.mkdir(parents=True, exist_ok=True)
with _DB_LOCK:
conn = get_conn()
try:
_create_schema(conn)
_apply_column_migrations(conn)
_ensure_email_unique_index(conn)
_migrate_user_emails(conn)
_ensure_initial_admin(conn)
conn.commit()
finally:
conn.close()
def get_secret_key() -> str:
"""获取(或生成)持久化的会话签名密钥(原子写入 + 失败显式告警)"""
key = ''
try:
if SECRET_KEY_PATH.exists():
key = SECRET_KEY_PATH.read_text(encoding='utf-8').strip()
if key:
return key
except OSError:
pass
new_key = secrets.token_hex(32)
try:
SECRET_KEY_PATH.parent.mkdir(parents=True, exist_ok=True)
# 先写临时文件再原子替换,避免部分写入留下损坏密钥
tmp = SECRET_KEY_PATH.with_suffix('.key.tmp')
tmp.write_text(new_key, encoding='utf-8')
os.replace(str(tmp), str(SECRET_KEY_PATH))
# 并发初始化者以"替换后文件里的值"为准(最后写入者胜出但全员一致)
final = SECRET_KEY_PATH.read_text(encoding='utf-8').strip()
return final or new_key
except OSError as e:
# 写入失败不能装作成功:显著告警,返回临时值(本进程会话可用,重启后失效)
try:
print(f'[警告] 会话密钥持久化失败({e}),本进程将使用临时密钥,重启后所有会话失效!', file=sys.stderr)
except UnicodeEncodeError:
print('[warn] secret key persistence failed; using temporary key for this process', file=sys.stderr)
return new_key