327 lines
13 KiB
Python
327 lines
13 KiB
Python
"""
|
||
公共模块：数据库连接、初始化、JWT 鉴权装饰器、模板渲染
|
||
供各蓝图复用
|
||
"""
|
||
import os
|
||
import secrets
|
||
from datetime import timedelta
|
||
from functools import wraps
|
||
|
||
import pymysql
|
||
from flask import request, redirect, url_for, jsonify, render_template, render_template_string, g
|
||
|
||
from config import mysql_host, mysql_user, mysql_password, mysql_database
|
||
from jwt_util import parse_token, get_token_from_request
|
||
|
||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
STATIC_DIR = os.path.join(BASE_DIR, 'static')
|
||
ASSETS_DIR = os.path.join(BASE_DIR, 'assets')
|
||
WEB_SOURCE_DIR = os.path.join(BASE_DIR, 'web_source')
|
||
TEMPLATE_FALLBACK_DIRS = (
|
||
WEB_SOURCE_DIR,
|
||
os.path.join(WEB_SOURCE_DIR, 'templates_backup'),
|
||
)
|
||
|
||
|
||
def get_db():
|
||
return pymysql.connect(
|
||
host=mysql_host,
|
||
user=mysql_user,
|
||
password=mysql_password,
|
||
database=mysql_database,
|
||
charset='utf8mb4',
|
||
cursorclass=pymysql.cursors.DictCursor
|
||
)
|
||
|
||
|
||
def _render_html(template_name: str, **context):
|
||
"""读取 HTML 模板：若为加密文件则先解密，再渲染。未加密或解密失败时按明文渲染。"""
|
||
path = next(
|
||
(candidate for candidate in (os.path.join(base_path, template_name) for base_path in TEMPLATE_FALLBACK_DIRS)
|
||
if os.path.isfile(candidate)),
|
||
None,
|
||
)
|
||
if path is None:
|
||
return render_template(template_name, **context)
|
||
with open(path, "rb") as f:
|
||
raw = f.read()
|
||
try:
|
||
from html_crypto import decrypt
|
||
content = decrypt(raw).decode("utf-8")
|
||
except Exception:
|
||
content = raw.decode("utf-8", errors="replace")
|
||
return render_template_string(content, **context)
|
||
|
||
|
||
def _render_html_new(template_name: str, **context):
|
||
"""读取 HTML 模板：若为加密文件则先解密，再渲染。未加密或解密失败时按明文渲染。"""
|
||
path = next(
|
||
(candidate for candidate in (os.path.join(base_path, template_name) for base_path in TEMPLATE_FALLBACK_DIRS)
|
||
if os.path.isfile(candidate)),
|
||
None,
|
||
)
|
||
if path is None:
|
||
return render_template(template_name, **context)
|
||
with open(path, "rb") as f:
|
||
raw = f.read()
|
||
try:
|
||
from html_crypto import decrypt
|
||
content = decrypt(raw).decode("utf-8")
|
||
except Exception:
|
||
content = raw.decode("utf-8", errors="replace")
|
||
return render_template_string(content, **context)
|
||
|
||
|
||
|
||
def _resolve_jwt_user():
|
||
"""从 Authorization/cookie 解析 JWT，命中后写入 flask.g 缓存。无效返回 None。"""
|
||
cached = getattr(g, "_aiimage_user", None)
|
||
if cached is not None:
|
||
return cached or None
|
||
token = get_token_from_request()
|
||
payload = parse_token(token) if token else None
|
||
if not payload:
|
||
g._aiimage_user = False
|
||
return None
|
||
g._aiimage_user = payload
|
||
g.aiimage_user_id = payload.get("user_id")
|
||
g.aiimage_username = payload.get("username") or ""
|
||
g.aiimage_device_id = payload.get("device_id") or ""
|
||
return payload
|
||
|
||
|
||
def current_user_id():
|
||
user = _resolve_jwt_user()
|
||
return user.get("user_id") if user else None
|
||
|
||
|
||
def current_username():
|
||
user = _resolve_jwt_user()
|
||
return user.get("username") if user else ""
|
||
|
||
|
||
def _is_session_user_valid():
|
||
"""兼容旧调用：判断当前 JWT 用户是否仍存在于数据库。"""
|
||
uid = current_user_id()
|
||
if not uid:
|
||
return False
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT id FROM users WHERE id = %s", (uid,))
|
||
row = cur.fetchone()
|
||
conn.close()
|
||
return bool(row)
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def _get_current_admin_role():
|
||
"""获取当前登录用户的管理角色：super_admin / admin / None（非管理员）"""
|
||
uid = current_user_id()
|
||
if not uid:
|
||
return None, None
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"SELECT id, username, is_admin, role, created_by_id FROM users WHERE id = %s",
|
||
(uid,)
|
||
)
|
||
row = cur.fetchone()
|
||
conn.close()
|
||
if not row or not row.get('is_admin'):
|
||
return None, None
|
||
return row.get('role') or ('super_admin' if row.get('created_by_id') is None else 'admin'), row
|
||
except Exception:
|
||
return None, None
|
||
|
||
|
||
def _unauthorized_response():
|
||
if request.headers.get('X-Requested-With') == 'XMLHttpRequest' \
|
||
or request.path.startswith('/api/') \
|
||
or 'application/json' in (request.headers.get('Accept') or ''):
|
||
return jsonify({'success': False, 'error': '未登录'}), 401
|
||
return redirect(url_for('auth.login'))
|
||
|
||
|
||
def login_required(f):
|
||
@wraps(f)
|
||
def decorated(*args, **kwargs):
|
||
if not current_user_id():
|
||
return _unauthorized_response()
|
||
return f(*args, **kwargs)
|
||
return decorated
|
||
|
||
|
||
def admin_required(f):
|
||
@wraps(f)
|
||
def decorated(*args, **kwargs):
|
||
uid = current_user_id()
|
||
if not uid:
|
||
return _unauthorized_response()
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT is_admin, role FROM users WHERE id = %s", (uid,))
|
||
row = cur.fetchone()
|
||
conn.close()
|
||
if not row or not row.get('is_admin'):
|
||
if request.headers.get('X-Requested-With') == 'XMLHttpRequest' \
|
||
or request.path.startswith('/api/') \
|
||
or 'application/json' in (request.headers.get('Accept') or ''):
|
||
return jsonify({'success': False, 'error': '需要管理员权限'}), 403
|
||
return redirect(url_for('main.home'))
|
||
except Exception as e:
|
||
if request.headers.get('X-Requested-With') == 'XMLHttpRequest' \
|
||
or request.path.startswith('/api/'):
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
return redirect(url_for('main.home'))
|
||
return f(*args, **kwargs)
|
||
return decorated
|
||
|
||
|
||
def init_db():
|
||
"""初始化数据库表，若不存在则创建"""
|
||
from werkzeug.security import generate_password_hash
|
||
conn = pymysql.connect(
|
||
host=mysql_host,
|
||
user=mysql_user,
|
||
password=mysql_password,
|
||
charset='utf8mb4'
|
||
)
|
||
try:
|
||
with conn.cursor() as cur:
|
||
cur.execute(f"CREATE DATABASE IF NOT EXISTS `{mysql_database}` DEFAULT CHARSET utf8mb4")
|
||
cur.execute(f"USE `{mysql_database}`")
|
||
cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS users (
|
||
id INT AUTO_INCREMENT PRIMARY KEY,
|
||
username VARCHAR(64) NOT NULL UNIQUE,
|
||
password_hash VARCHAR(256) NOT NULL,
|
||
is_admin TINYINT(1) DEFAULT 0,
|
||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
||
)
|
||
""")
|
||
cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS image_history (
|
||
id INT AUTO_INCREMENT PRIMARY KEY,
|
||
user_id INT NOT NULL,
|
||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||
panel_type VARCHAR(64) DEFAULT '',
|
||
original_urls JSON,
|
||
params JSON,
|
||
result_urls JSON,
|
||
long_image_url VARCHAR(1024) DEFAULT NULL,
|
||
INDEX idx_user_created (user_id, created_at DESC)
|
||
)
|
||
""")
|
||
try:
|
||
cur.execute("ALTER TABLE image_history ADD COLUMN long_image_url VARCHAR(1024) DEFAULT NULL")
|
||
except Exception:
|
||
pass
|
||
try:
|
||
cur.execute("ALTER TABLE users ADD COLUMN machine VARCHAR(64) DEFAULT NULL")
|
||
except Exception:
|
||
pass
|
||
try:
|
||
cur.execute("ALTER TABLE users ADD COLUMN role VARCHAR(20) DEFAULT 'normal'")
|
||
except Exception:
|
||
pass
|
||
try:
|
||
cur.execute("ALTER TABLE users ADD COLUMN created_by_id INT NULL")
|
||
except Exception:
|
||
pass
|
||
cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS brand_crawl_tasks (
|
||
id INT AUTO_INCREMENT PRIMARY KEY,
|
||
user_id INT NOT NULL,
|
||
file_paths JSON NOT NULL,
|
||
status VARCHAR(20) DEFAULT 'pending',
|
||
task_type TINYINT DEFAULT 1 COMMENT '1=立即执行,2=添加任务执行',
|
||
result_paths JSON NULL,
|
||
error_message TEXT NULL,
|
||
progress_current INT DEFAULT 0,
|
||
progress_total INT DEFAULT 0,
|
||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
|
||
INDEX idx_user_status (user_id, status)
|
||
)
|
||
""")
|
||
try:
|
||
cur.execute("ALTER TABLE brand_crawl_tasks ADD COLUMN progress_current INT DEFAULT 0")
|
||
except Exception:
|
||
pass
|
||
try:
|
||
cur.execute("ALTER TABLE brand_crawl_tasks ADD COLUMN progress_total INT DEFAULT 0")
|
||
except Exception:
|
||
pass
|
||
try:
|
||
cur.execute("ALTER TABLE brand_crawl_tasks ADD COLUMN `desc` VARCHAR(500) NULL COMMENT '上传文件名描述'")
|
||
except Exception:
|
||
pass
|
||
try:
|
||
cur.execute("ALTER TABLE brand_crawl_tasks ADD COLUMN strategy VARCHAR(20) DEFAULT 'Terms' COMMENT '品牌匹配方式: Terms/Simple'")
|
||
except Exception:
|
||
pass
|
||
try:
|
||
cur.execute("ALTER TABLE brand_crawl_tasks ADD COLUMN task_type TINYINT DEFAULT 1 COMMENT '1=立即执行,2=添加任务执行'")
|
||
except Exception:
|
||
pass
|
||
cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS columns (
|
||
id INT AUTO_INCREMENT PRIMARY KEY,
|
||
name VARCHAR(128) NOT NULL COMMENT '栏目名',
|
||
column_key VARCHAR(64) NOT NULL COMMENT '栏目标识',
|
||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||
UNIQUE KEY uk_column_key (column_key)
|
||
)
|
||
""")
|
||
cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS user_column_permission (
|
||
user_id INT NOT NULL,
|
||
column_id INT NOT NULL,
|
||
PRIMARY KEY (user_id, column_id),
|
||
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE,
|
||
FOREIGN KEY (column_id) REFERENCES columns(id) ON DELETE CASCADE
|
||
)
|
||
""")
|
||
try:
|
||
cur.execute("UPDATE users SET role = 'normal' WHERE (role IS NULL OR role = '') AND (is_admin = 0 OR is_admin IS NULL)")
|
||
cur.execute("SELECT MIN(id) AS mid FROM users WHERE is_admin = 1")
|
||
row = cur.fetchone()
|
||
if row and row.get('mid'):
|
||
mid = row['mid']
|
||
cur.execute("UPDATE users SET role = 'super_admin' WHERE id = %s", (mid,))
|
||
cur.execute("UPDATE users SET role = 'admin' WHERE is_admin = 1 AND id != %s", (mid,))
|
||
cur.execute("UPDATE users SET created_by_id = %s WHERE role = 'admin' AND (created_by_id IS NULL)", (mid,))
|
||
except Exception:
|
||
pass
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
_create_initial_admin()
|
||
|
||
|
||
def _create_initial_admin():
|
||
"""若没有任何管理员，则创建默认超级管理员（首次启动时，仅一个）"""
|
||
from werkzeug.security import generate_password_hash
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT id FROM users WHERE role = 'super_admin' LIMIT 1")
|
||
if cur.fetchone():
|
||
conn.close()
|
||
return
|
||
admin_user = os.environ.get('ADMIN_USER', 'admin')
|
||
admin_pwd = os.environ.get('ADMIN_PASSWORD', 'admin123')
|
||
pwd_hash = generate_password_hash(admin_pwd, method='pbkdf2:sha256')
|
||
cur.execute(
|
||
"INSERT INTO users (username, password_hash, is_admin, role) VALUES (%s, %s, 1, 'super_admin')",
|
||
(admin_user, pwd_hash)
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
except Exception:
|
||
pass
|