428 lines
14 KiB
Python
428 lines
14 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import sqlite3
|
||
from pathlib import Path
|
||
|
||
import db_pool
|
||
from db.utils import get_cst_now_str
|
||
from password_utils import (
|
||
hash_password_bcrypt,
|
||
is_sha256_hash,
|
||
verify_password_bcrypt,
|
||
verify_password_sha256,
|
||
)
|
||
|
||
_DEFAULT_SYSTEM_CONFIG = {
|
||
"max_concurrent_global": 2,
|
||
"max_concurrent_per_account": 1,
|
||
"max_screenshot_concurrent": 3,
|
||
"db_slow_query_ms": 120,
|
||
"schedule_enabled": 0,
|
||
"schedule_time": "02:00",
|
||
"schedule_browse_type": "应读",
|
||
"schedule_weekdays": "1,2,3,4,5,6,7",
|
||
"proxy_enabled": 0,
|
||
"proxy_api_url": "",
|
||
"proxy_expire_minutes": 3,
|
||
"enable_screenshot": 1,
|
||
"auto_approve_enabled": 0,
|
||
"auto_approve_hourly_limit": 10,
|
||
"auto_approve_vip_days": 7,
|
||
"kdocs_enabled": 0,
|
||
"kdocs_doc_url": "",
|
||
"kdocs_default_unit": "",
|
||
"kdocs_sheet_name": "",
|
||
"kdocs_sheet_index": 0,
|
||
"kdocs_unit_column": "A",
|
||
"kdocs_image_column": "D",
|
||
"kdocs_admin_notify_enabled": 0,
|
||
"kdocs_admin_notify_email": "",
|
||
"kdocs_row_start": 0,
|
||
"kdocs_row_end": 0,
|
||
}
|
||
|
||
_SYSTEM_CONFIG_UPDATERS = (
|
||
("max_concurrent_global", "max_concurrent"),
|
||
("schedule_enabled", "schedule_enabled"),
|
||
("schedule_time", "schedule_time"),
|
||
("schedule_browse_type", "schedule_browse_type"),
|
||
("schedule_weekdays", "schedule_weekdays"),
|
||
("max_concurrent_per_account", "max_concurrent_per_account"),
|
||
("max_screenshot_concurrent", "max_screenshot_concurrent"),
|
||
("db_slow_query_ms", "db_slow_query_ms"),
|
||
("enable_screenshot", "enable_screenshot"),
|
||
("proxy_enabled", "proxy_enabled"),
|
||
("proxy_api_url", "proxy_api_url"),
|
||
("proxy_expire_minutes", "proxy_expire_minutes"),
|
||
("auto_approve_enabled", "auto_approve_enabled"),
|
||
("auto_approve_hourly_limit", "auto_approve_hourly_limit"),
|
||
("auto_approve_vip_days", "auto_approve_vip_days"),
|
||
("kdocs_enabled", "kdocs_enabled"),
|
||
("kdocs_doc_url", "kdocs_doc_url"),
|
||
("kdocs_default_unit", "kdocs_default_unit"),
|
||
("kdocs_sheet_name", "kdocs_sheet_name"),
|
||
("kdocs_sheet_index", "kdocs_sheet_index"),
|
||
("kdocs_unit_column", "kdocs_unit_column"),
|
||
("kdocs_image_column", "kdocs_image_column"),
|
||
("kdocs_admin_notify_enabled", "kdocs_admin_notify_enabled"),
|
||
("kdocs_admin_notify_email", "kdocs_admin_notify_email"),
|
||
("kdocs_row_start", "kdocs_row_start"),
|
||
("kdocs_row_end", "kdocs_row_end"),
|
||
)
|
||
|
||
|
||
def _count_scalar(cursor, sql: str, params=()) -> int:
|
||
cursor.execute(sql, params)
|
||
row = cursor.fetchone()
|
||
if not row:
|
||
return 0
|
||
try:
|
||
if "count" in row.keys():
|
||
return int(row["count"] or 0)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
return int(row[0] or 0)
|
||
except Exception:
|
||
return 0
|
||
|
||
|
||
def _table_exists(cursor, table_name: str) -> bool:
|
||
cursor.execute(
|
||
"""
|
||
SELECT name FROM sqlite_master
|
||
WHERE type='table' AND name=?
|
||
""",
|
||
(table_name,),
|
||
)
|
||
return bool(cursor.fetchone())
|
||
|
||
|
||
def _normalize_days(days, default: int = 30) -> int:
|
||
try:
|
||
value = int(days)
|
||
except Exception:
|
||
value = default
|
||
if value < 0:
|
||
return 0
|
||
return value
|
||
|
||
|
||
def _store_default_admin_credentials(username: str, password: str) -> str | None:
|
||
"""将首次管理员账号密码写入受限权限文件，避免打印到日志。"""
|
||
raw_path = str(
|
||
os.environ.get("DEFAULT_ADMIN_CREDENTIALS_FILE", "data/default_admin_credentials.txt") or ""
|
||
).strip()
|
||
if not raw_path:
|
||
return None
|
||
|
||
cred_path = Path(raw_path)
|
||
try:
|
||
cred_path.parent.mkdir(parents=True, exist_ok=True)
|
||
with open(cred_path, "w", encoding="utf-8") as f:
|
||
f.write("安全提醒：首次管理员账号已创建\n")
|
||
f.write(f"用户名: {username}\n")
|
||
f.write(f"密码: {password}\n")
|
||
f.write("请登录后立即修改密码，并删除该文件。\n")
|
||
os.chmod(cred_path, 0o600)
|
||
return str(cred_path)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def ensure_default_admin() -> bool:
|
||
"""确保存在默认管理员账号（行为保持不变）。"""
|
||
import secrets
|
||
import string
|
||
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
count = _count_scalar(cursor, "SELECT COUNT(*) as count FROM admins")
|
||
|
||
if count == 0:
|
||
alphabet = string.ascii_letters + string.digits
|
||
bootstrap_password = str(os.environ.get("DEFAULT_ADMIN_PASSWORD", "") or "").strip()
|
||
random_password = bootstrap_password or "".join(secrets.choice(alphabet) for _ in range(12))
|
||
|
||
default_password_hash = hash_password_bcrypt(random_password)
|
||
cursor.execute(
|
||
"INSERT INTO admins (username, password_hash, created_at) VALUES (?, ?, ?)",
|
||
("admin", default_password_hash, get_cst_now_str()),
|
||
)
|
||
conn.commit()
|
||
credential_file = _store_default_admin_credentials("admin", random_password)
|
||
print("=" * 60)
|
||
print("安全提醒：已创建默认管理员账号")
|
||
print("用户名: admin")
|
||
if credential_file:
|
||
print(f"初始密码已写入: {credential_file}（权限600）")
|
||
print("请立即登录后修改密码，并删除该文件。")
|
||
else:
|
||
print("未能写入初始密码文件。")
|
||
print("建议设置 DEFAULT_ADMIN_PASSWORD 后重建管理员账号。")
|
||
print("=" * 60)
|
||
return True
|
||
return False
|
||
|
||
|
||
def verify_admin(username: str, password: str):
|
||
"""验证管理员登录 - 自动从SHA256升级到bcrypt（行为保持不变）。"""
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
cursor.execute("SELECT * FROM admins WHERE username = ?", (username,))
|
||
admin = cursor.fetchone()
|
||
|
||
if not admin:
|
||
return None
|
||
|
||
admin_dict = dict(admin)
|
||
password_hash = admin_dict["password_hash"]
|
||
|
||
if is_sha256_hash(password_hash):
|
||
if verify_password_sha256(password, password_hash):
|
||
new_hash = hash_password_bcrypt(password)
|
||
cursor.execute("UPDATE admins SET password_hash = ? WHERE username = ?", (new_hash, username))
|
||
conn.commit()
|
||
print(f"管理员 {username} 密码已自动升级到bcrypt")
|
||
return admin_dict
|
||
return None
|
||
|
||
if verify_password_bcrypt(password, password_hash):
|
||
return admin_dict
|
||
return None
|
||
|
||
|
||
def get_admin_by_username(username: str):
|
||
"""根据用户名获取管理员记录"""
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
cursor.execute("SELECT * FROM admins WHERE username = ?", (username,))
|
||
row = cursor.fetchone()
|
||
return dict(row) if row else None
|
||
|
||
|
||
def get_admin_by_id(admin_id: int):
|
||
"""根据ID获取管理员记录"""
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
cursor.execute("SELECT * FROM admins WHERE id = ?", (int(admin_id),))
|
||
row = cursor.fetchone()
|
||
return dict(row) if row else None
|
||
|
||
|
||
def update_admin_password(username: str, new_password: str) -> bool:
|
||
"""更新管理员密码"""
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
password_hash = hash_password_bcrypt(new_password)
|
||
cursor.execute("UPDATE admins SET password_hash = ? WHERE username = ?", (password_hash, username))
|
||
conn.commit()
|
||
return cursor.rowcount > 0
|
||
|
||
|
||
def update_admin_username(old_username: str, new_username: str) -> bool:
|
||
"""更新管理员用户名"""
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
try:
|
||
cursor.execute("UPDATE admins SET username = ? WHERE username = ?", (new_username, old_username))
|
||
conn.commit()
|
||
return True
|
||
except sqlite3.IntegrityError:
|
||
return False
|
||
|
||
|
||
def get_system_stats() -> dict:
|
||
"""获取系统统计信息"""
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
|
||
cursor.execute(
|
||
"""
|
||
SELECT
|
||
COUNT(*) AS total_users,
|
||
SUM(CASE WHEN status = 'approved' THEN 1 ELSE 0 END) AS approved_users,
|
||
SUM(CASE WHEN date(created_at) = date('now', 'localtime') THEN 1 ELSE 0 END) AS new_users_today,
|
||
SUM(CASE WHEN datetime(created_at) >= datetime('now', 'localtime', '-7 days') THEN 1 ELSE 0 END) AS new_users_7d,
|
||
SUM(
|
||
CASE
|
||
WHEN vip_expire_time IS NOT NULL
|
||
AND datetime(vip_expire_time) > datetime('now', 'localtime')
|
||
THEN 1 ELSE 0
|
||
END
|
||
) AS vip_users
|
||
FROM users
|
||
"""
|
||
)
|
||
user_stats = cursor.fetchone() or {}
|
||
|
||
def _to_int(key: str) -> int:
|
||
try:
|
||
return int(user_stats[key] or 0)
|
||
except Exception:
|
||
return 0
|
||
|
||
total_accounts = _count_scalar(cursor, "SELECT COUNT(*) as count FROM accounts")
|
||
|
||
return {
|
||
"total_users": _to_int("total_users"),
|
||
"approved_users": _to_int("approved_users"),
|
||
"new_users_today": _to_int("new_users_today"),
|
||
"new_users_7d": _to_int("new_users_7d"),
|
||
"total_accounts": total_accounts,
|
||
"vip_users": _to_int("vip_users"),
|
||
}
|
||
|
||
|
||
def get_system_config_raw() -> dict:
|
||
"""获取系统配置（无缓存，供 facade 做缓存/失效）。"""
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
cursor.execute("SELECT * FROM system_config WHERE id = 1")
|
||
row = cursor.fetchone()
|
||
if row:
|
||
return dict(row)
|
||
return dict(_DEFAULT_SYSTEM_CONFIG)
|
||
|
||
|
||
def update_system_config(
|
||
*,
|
||
max_concurrent=None,
|
||
schedule_enabled=None,
|
||
schedule_time=None,
|
||
schedule_browse_type=None,
|
||
schedule_weekdays=None,
|
||
max_concurrent_per_account=None,
|
||
max_screenshot_concurrent=None,
|
||
enable_screenshot=None,
|
||
proxy_enabled=None,
|
||
proxy_api_url=None,
|
||
proxy_expire_minutes=None,
|
||
auto_approve_enabled=None,
|
||
auto_approve_hourly_limit=None,
|
||
auto_approve_vip_days=None,
|
||
kdocs_enabled=None,
|
||
kdocs_doc_url=None,
|
||
kdocs_default_unit=None,
|
||
kdocs_sheet_name=None,
|
||
kdocs_sheet_index=None,
|
||
kdocs_unit_column=None,
|
||
kdocs_image_column=None,
|
||
kdocs_admin_notify_enabled=None,
|
||
kdocs_admin_notify_email=None,
|
||
kdocs_row_start=None,
|
||
kdocs_row_end=None,
|
||
db_slow_query_ms=None,
|
||
) -> bool:
|
||
"""更新系统配置（仅更新DB，不做缓存处理）。"""
|
||
arg_values = {
|
||
"max_concurrent": max_concurrent,
|
||
"schedule_enabled": schedule_enabled,
|
||
"schedule_time": schedule_time,
|
||
"schedule_browse_type": schedule_browse_type,
|
||
"schedule_weekdays": schedule_weekdays,
|
||
"max_concurrent_per_account": max_concurrent_per_account,
|
||
"max_screenshot_concurrent": max_screenshot_concurrent,
|
||
"enable_screenshot": enable_screenshot,
|
||
"proxy_enabled": proxy_enabled,
|
||
"proxy_api_url": proxy_api_url,
|
||
"proxy_expire_minutes": proxy_expire_minutes,
|
||
"auto_approve_enabled": auto_approve_enabled,
|
||
"auto_approve_hourly_limit": auto_approve_hourly_limit,
|
||
"auto_approve_vip_days": auto_approve_vip_days,
|
||
"kdocs_enabled": kdocs_enabled,
|
||
"kdocs_doc_url": kdocs_doc_url,
|
||
"kdocs_default_unit": kdocs_default_unit,
|
||
"kdocs_sheet_name": kdocs_sheet_name,
|
||
"kdocs_sheet_index": kdocs_sheet_index,
|
||
"kdocs_unit_column": kdocs_unit_column,
|
||
"kdocs_image_column": kdocs_image_column,
|
||
"kdocs_admin_notify_enabled": kdocs_admin_notify_enabled,
|
||
"kdocs_admin_notify_email": kdocs_admin_notify_email,
|
||
"kdocs_row_start": kdocs_row_start,
|
||
"kdocs_row_end": kdocs_row_end,
|
||
"db_slow_query_ms": db_slow_query_ms,
|
||
}
|
||
|
||
updates = []
|
||
params = []
|
||
for db_field, arg_name in _SYSTEM_CONFIG_UPDATERS:
|
||
value = arg_values.get(arg_name)
|
||
if value is None:
|
||
continue
|
||
updates.append(f"{db_field} = ?")
|
||
params.append(value)
|
||
|
||
if not updates:
|
||
return False
|
||
|
||
updates.append("updated_at = ?")
|
||
params.append(get_cst_now_str())
|
||
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
sql = f"UPDATE system_config SET {', '.join(updates)} WHERE id = 1"
|
||
cursor.execute(sql, params)
|
||
conn.commit()
|
||
return True
|
||
|
||
|
||
def get_hourly_registration_count() -> int:
|
||
"""获取最近一小时内的注册用户数"""
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
return _count_scalar(
|
||
cursor,
|
||
"""
|
||
SELECT COUNT(*) as count FROM users
|
||
WHERE created_at >= datetime('now', 'localtime', '-1 hour')
|
||
""",
|
||
)
|
||
|
||
|
||
# ==================== 密码重置（管理员） ====================
|
||
|
||
|
||
def admin_reset_user_password(user_id: int, new_password: str) -> bool:
|
||
"""管理员直接重置用户密码"""
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
password_hash = hash_password_bcrypt(new_password)
|
||
try:
|
||
cursor.execute("UPDATE users SET password_hash = ? WHERE id = ?", (password_hash, user_id))
|
||
conn.commit()
|
||
return cursor.rowcount > 0
|
||
except Exception as e:
|
||
print(f"管理员重置密码失败: {e}")
|
||
return False
|
||
|
||
|
||
def clean_old_operation_logs(days: int = 30) -> int:
|
||
"""清理指定天数前的操作日志（如果存在operation_logs表）"""
|
||
safe_days = _normalize_days(days, default=30)
|
||
|
||
with db_pool.get_db() as conn:
|
||
cursor = conn.cursor()
|
||
|
||
if not _table_exists(cursor, "operation_logs"):
|
||
return 0
|
||
|
||
try:
|
||
cursor.execute(
|
||
"""
|
||
DELETE FROM operation_logs
|
||
WHERE created_at < datetime('now', 'localtime', '-' || ? || ' days')
|
||
""",
|
||
(safe_days,),
|
||
)
|
||
deleted_count = cursor.rowcount
|
||
conn.commit()
|
||
print(f"已清理 {deleted_count} 条旧操作日志 (>{safe_days}天)")
|
||
return deleted_count
|
||
except Exception as e:
|
||
print(f"清理旧操作日志失败: {e}")
|
||
return 0
|