373 lines
13 KiB
Python
373 lines
13 KiB
Python
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import hmac
|
||
import secrets
|
||
import sqlite3
|
||
import threading
|
||
import time
|
||
from contextlib import contextmanager
|
||
from dataclasses import dataclass
|
||
from datetime import datetime, timezone
|
||
from http.cookies import SimpleCookie
|
||
from typing import Iterator
|
||
|
||
from .config import (
|
||
AUTH_COOKIE_SECURE,
|
||
AUTH_DB_PATH,
|
||
AUTH_SESSION_COOKIE_NAME,
|
||
AUTH_SESSION_TTL_HOURS,
|
||
utc_now,
|
||
)
|
||
|
||
# 默认账号只在服务端本地初始化，避免出现在 README 或 .env 模板里。
|
||
BOOTSTRAP_USERNAME = "haiyi"
|
||
# 初始密码固定按需求内置，登录时由用户在页面手动输入。
|
||
BOOTSTRAP_PASSWORD = "haiyi@1234.com"
|
||
# 密码统一走 PBKDF2，避免把明文写进本地数据库。
|
||
PASSWORD_HASH_ALGORITHM = "pbkdf2_sha256"
|
||
# 迭代次数拉高一些，降低弱口令被离线撞库的风险。
|
||
PASSWORD_HASH_ITERATIONS = 600_000
|
||
# Session token 使用 32 字节随机数，足够覆盖当前单机面板场景。
|
||
SESSION_TOKEN_BYTES = 32
|
||
# 初始化表结构时加锁，避免线程服务器并发首访时重复建表。
|
||
INIT_LOCK = threading.Lock()
|
||
# 进程内记一个初始化标记，减少每次请求都跑 schema 检查。
|
||
INITIALIZED = False
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class SessionRecord:
|
||
# 当前登录用户名。
|
||
username: str
|
||
# 会话过期时间，返回给前端展示或调试。
|
||
expires_at: str
|
||
# 原始时间戳，服务端判断是否过期时直接复用。
|
||
expires_at_ts: int
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class CreatedSession(SessionRecord):
|
||
# 登录成功后下发到 Cookie 的随机 token。
|
||
token: str
|
||
|
||
|
||
def ensure_auth_store() -> None:
|
||
global INITIALIZED
|
||
|
||
# 已初始化过且文件还在时，直接复用现有状态。
|
||
if INITIALIZED and AUTH_DB_PATH.exists():
|
||
return
|
||
|
||
with INIT_LOCK:
|
||
# 双重检查，避免多个线程排队后重复建表。
|
||
if INITIALIZED and AUTH_DB_PATH.exists():
|
||
return
|
||
|
||
# 先确保数据库目录存在，SQLite 才能正常创建文件。
|
||
AUTH_DB_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
with db_connection() as connection:
|
||
# 单机线程服务读多写少，WAL 更适合并发读场景。
|
||
connection.execute("PRAGMA journal_mode = WAL")
|
||
# 用户表只保存账号和密码摘要。
|
||
connection.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS users (
|
||
username TEXT PRIMARY KEY,
|
||
password_hash TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
updated_at TEXT NOT NULL,
|
||
last_login_at TEXT NOT NULL DEFAULT ''
|
||
)
|
||
"""
|
||
)
|
||
# 会话表按 token hash 存储，数据库泄漏时也不直接暴露真实 token。
|
||
connection.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS sessions (
|
||
token_hash TEXT PRIMARY KEY,
|
||
username TEXT NOT NULL,
|
||
created_at TEXT NOT NULL,
|
||
expires_at INTEGER NOT NULL,
|
||
remote_addr TEXT NOT NULL DEFAULT '',
|
||
user_agent TEXT NOT NULL DEFAULT '',
|
||
FOREIGN KEY (username) REFERENCES users(username) ON DELETE CASCADE
|
||
)
|
||
"""
|
||
)
|
||
# 过期清理主要按 expires_at 查找，这里单独建索引。
|
||
connection.execute(
|
||
"""
|
||
CREATE INDEX IF NOT EXISTS idx_sessions_expires_at
|
||
ON sessions(expires_at)
|
||
"""
|
||
)
|
||
|
||
# 默认账号只在缺失时自动补，避免覆盖已有密码。
|
||
bootstrap_user = connection.execute(
|
||
"SELECT username FROM users WHERE username = ?",
|
||
(BOOTSTRAP_USERNAME,),
|
||
).fetchone()
|
||
if bootstrap_user is None:
|
||
now = utc_now()
|
||
connection.execute(
|
||
"""
|
||
INSERT INTO users(username, password_hash, created_at, updated_at, last_login_at)
|
||
VALUES (?, ?, ?, ?, '')
|
||
""",
|
||
(
|
||
BOOTSTRAP_USERNAME,
|
||
hash_password(BOOTSTRAP_PASSWORD),
|
||
now,
|
||
now,
|
||
),
|
||
)
|
||
|
||
# 只有建表和默认账号都准备好后，才把进程状态切成已初始化。
|
||
INITIALIZED = True
|
||
|
||
|
||
def connect_db() -> sqlite3.Connection:
|
||
# SQLite 每次请求单独建连接，避免线程间复用同一连接。
|
||
connection = sqlite3.connect(AUTH_DB_PATH, timeout=5.0)
|
||
# 结果行按列名读取，后续代码更直观。
|
||
connection.row_factory = sqlite3.Row
|
||
return connection
|
||
|
||
|
||
@contextmanager
|
||
def db_connection() -> Iterator[sqlite3.Connection]:
|
||
# sqlite3.Connection 的 with 语法不会关闭连接，统一在这里提交并关闭。
|
||
connection = connect_db()
|
||
try:
|
||
yield connection
|
||
connection.commit()
|
||
except Exception:
|
||
connection.rollback()
|
||
raise
|
||
finally:
|
||
connection.close()
|
||
|
||
|
||
def hash_password(password: str, *, salt: str | None = None) -> str:
|
||
# 没传 salt 时生成新的随机盐值。
|
||
actual_salt = salt or secrets.token_hex(16)
|
||
# PBKDF2 输出二进制摘要，最后统一转十六进制文本保存。
|
||
digest = hashlib.pbkdf2_hmac(
|
||
"sha256",
|
||
password.encode("utf-8"),
|
||
actual_salt.encode("utf-8"),
|
||
PASSWORD_HASH_ITERATIONS,
|
||
).hex()
|
||
# 统一用固定格式拼接，校验时便于解析算法、轮数和盐值。
|
||
return f"{PASSWORD_HASH_ALGORITHM}${PASSWORD_HASH_ITERATIONS}${actual_salt}${digest}"
|
||
|
||
|
||
def verify_password(password: str, stored_hash: str) -> bool:
|
||
try:
|
||
# 历史格式不正确时直接按失败处理，避免抛内部异常到接口。
|
||
algorithm, raw_iterations, salt, expected_digest = stored_hash.split("$", 3)
|
||
except ValueError:
|
||
return False
|
||
|
||
# 当前版本只接受约定的 PBKDF2 格式。
|
||
if algorithm != PASSWORD_HASH_ALGORITHM:
|
||
return False
|
||
|
||
try:
|
||
# 迭代次数也从存储值里解析，后续如果升级策略兼容性更好。
|
||
iterations = int(raw_iterations)
|
||
except ValueError:
|
||
return False
|
||
|
||
actual_digest = hashlib.pbkdf2_hmac(
|
||
"sha256",
|
||
password.encode("utf-8"),
|
||
salt.encode("utf-8"),
|
||
iterations,
|
||
).hex()
|
||
# 结果比较走常量时间比对，减少时序侧信道。
|
||
return hmac.compare_digest(actual_digest, expected_digest)
|
||
|
||
|
||
def authenticate_user(username: str, password: str) -> str | None:
|
||
ensure_auth_store()
|
||
|
||
# 空账号或空密码都直接拒绝，避免继续查库。
|
||
if not username or not password:
|
||
return None
|
||
|
||
with db_connection() as connection:
|
||
# 当前账号只支持精确匹配，避免大小写模糊带来歧义。
|
||
row = connection.execute(
|
||
"SELECT username, password_hash FROM users WHERE username = ?",
|
||
(username,),
|
||
).fetchone()
|
||
|
||
# 用户不存在时返回空，后续统一走登录失败。
|
||
if row is None:
|
||
return None
|
||
|
||
# 密码摘要校验失败也返回空，不向外暴露更多信息。
|
||
if not verify_password(password, str(row["password_hash"])):
|
||
return None
|
||
|
||
# 认证成功时返回规范用户名。
|
||
return str(row["username"])
|
||
|
||
|
||
def create_session(username: str, *, remote_addr: str = "", user_agent: str = "") -> CreatedSession:
|
||
ensure_auth_store()
|
||
|
||
# 登录成功后立刻生成新的随机 token，避免复用旧会话。
|
||
token = secrets.token_urlsafe(SESSION_TOKEN_BYTES)
|
||
# 真实 token 只放到 Cookie，库里只保留它的哈希值。
|
||
token_hash = hash_token(token)
|
||
# 过期时间直接按小时配置换算成秒。
|
||
expires_at_ts = int(time.time()) + max(1, AUTH_SESSION_TTL_HOURS) * 3600
|
||
now = utc_now()
|
||
|
||
with db_connection() as connection:
|
||
# 每次登录前顺手清理一遍过期会话，避免本地库无限增长。
|
||
purge_expired_sessions(connection=connection)
|
||
# 新会话和来源信息一起落库，后续排查时能知道来源。
|
||
connection.execute(
|
||
"""
|
||
INSERT INTO sessions(token_hash, username, created_at, expires_at, remote_addr, user_agent)
|
||
VALUES (?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(token_hash, username, now, expires_at_ts, remote_addr, user_agent),
|
||
)
|
||
# 登录成功时更新最后登录时间。
|
||
connection.execute(
|
||
"UPDATE users SET last_login_at = ? WHERE username = ?",
|
||
(now, username),
|
||
)
|
||
|
||
# 把可直接写回前端的会话信息一起返回。
|
||
return CreatedSession(
|
||
username=username,
|
||
expires_at=format_timestamp(expires_at_ts),
|
||
expires_at_ts=expires_at_ts,
|
||
token=token,
|
||
)
|
||
|
||
|
||
def get_session(token: str) -> SessionRecord | None:
|
||
ensure_auth_store()
|
||
|
||
# 没有 token 时直接视为未登录。
|
||
if not token:
|
||
return None
|
||
|
||
now_ts = int(time.time())
|
||
token_hash = hash_token(token)
|
||
|
||
with db_connection() as connection:
|
||
# 先删掉已过期的历史会话，减少后续查询噪音。
|
||
purge_expired_sessions(connection=connection, now_ts=now_ts)
|
||
row = connection.execute(
|
||
"""
|
||
SELECT username, expires_at
|
||
FROM sessions
|
||
WHERE token_hash = ? AND expires_at > ?
|
||
""",
|
||
(token_hash, now_ts),
|
||
).fetchone()
|
||
|
||
# 没查到时说明 token 不存在或已经过期。
|
||
if row is None:
|
||
return None
|
||
|
||
expires_at_ts = int(row["expires_at"])
|
||
return SessionRecord(
|
||
username=str(row["username"]),
|
||
expires_at=format_timestamp(expires_at_ts),
|
||
expires_at_ts=expires_at_ts,
|
||
)
|
||
|
||
|
||
def delete_session(token: str) -> None:
|
||
ensure_auth_store()
|
||
|
||
# 空 token 不需要继续删库。
|
||
if not token:
|
||
return
|
||
|
||
with db_connection() as connection:
|
||
# 删除不存在的 token 也视为成功，便于前端幂等退出。
|
||
connection.execute("DELETE FROM sessions WHERE token_hash = ?", (hash_token(token),))
|
||
|
||
|
||
def purge_expired_sessions(*, connection: sqlite3.Connection | None = None, now_ts: int | None = None) -> None:
|
||
# 没传时间时按当前时间清理。
|
||
actual_now_ts = now_ts if now_ts is not None else int(time.time())
|
||
|
||
# 外部没传连接时自行开一个短连接完成清理。
|
||
if connection is None:
|
||
with db_connection() as owned_connection:
|
||
owned_connection.execute("DELETE FROM sessions WHERE expires_at <= ?", (actual_now_ts,))
|
||
return
|
||
|
||
# 复用外部事务时只执行删除，不在这里额外提交。
|
||
connection.execute("DELETE FROM sessions WHERE expires_at <= ?", (actual_now_ts,))
|
||
|
||
|
||
def hash_token(token: str) -> str:
|
||
# token 哈希统一走 SHA-256，足够覆盖当前随机 token 长度。
|
||
return hashlib.sha256(token.encode("utf-8")).hexdigest()
|
||
|
||
|
||
def parse_session_token(cookie_header: str | None) -> str:
|
||
# 没 Cookie 头时直接返回空。
|
||
if not cookie_header:
|
||
return ""
|
||
|
||
cookie = SimpleCookie()
|
||
try:
|
||
# 非法 Cookie 头统一按未登录处理。
|
||
cookie.load(cookie_header)
|
||
except Exception: # noqa: BLE001
|
||
return ""
|
||
|
||
morsel = cookie.get(AUTH_SESSION_COOKIE_NAME)
|
||
# 没找到目标 Cookie 时同样返回空串。
|
||
return morsel.value if morsel is not None else ""
|
||
|
||
|
||
def build_session_cookie(token: str, *, max_age: int) -> str:
|
||
# 登录成功后写入新的 HttpOnly Cookie。
|
||
return build_cookie_header(token, max_age=max_age)
|
||
|
||
|
||
def build_clear_session_cookie() -> str:
|
||
# 退出登录或 token 无效时，把 Cookie 直接清空。
|
||
return build_cookie_header("", max_age=0, expires="Thu, 01 Jan 1970 00:00:00 GMT")
|
||
|
||
|
||
def build_cookie_header(value: str, *, max_age: int, expires: str | None = None) -> str:
|
||
cookie = SimpleCookie()
|
||
cookie[AUTH_SESSION_COOKIE_NAME] = value
|
||
morsel = cookie[AUTH_SESSION_COOKIE_NAME]
|
||
# 所有页面都复用同一套登录态。
|
||
morsel["path"] = "/"
|
||
# 登录态不允许被前端脚本直接读取。
|
||
morsel["httponly"] = True
|
||
# 面板是同域使用，Strict 可以把 CSRF 风险压低。
|
||
morsel["samesite"] = "Strict"
|
||
# 当前仍有 HTTP 直连场景，因此只在显式配置时才打开 Secure。
|
||
if AUTH_COOKIE_SECURE:
|
||
morsel["secure"] = True
|
||
# max-age 统一写秒数，浏览器过期处理更直接。
|
||
morsel["max-age"] = str(max_age)
|
||
# 需要立即清空时同时补 expires，兼容性更稳。
|
||
if expires:
|
||
morsel["expires"] = expires
|
||
return morsel.OutputString()
|
||
|
||
|
||
def format_timestamp(timestamp: int) -> str:
|
||
# Session 返回给前端时统一输出 UTC ISO 时间。
|
||
return datetime.fromtimestamp(timestamp, timezone.utc).replace(microsecond=0).isoformat()
|