""" 公共模块:数据库连接、初始化、会话校验、装饰器、模板渲染 供各蓝图复用 """ import os import secrets from datetime import timedelta from functools import wraps import pymysql from flask import request, redirect, url_for, session, jsonify, render_template, render_template_string from config import mysql_host, mysql_user, mysql_password, mysql_database BASE_DIR = os.path.dirname(os.path.abspath(__file__)) STATIC_DIR = os.path.join(BASE_DIR, 'static') 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 = os.path.join(BASE_DIR, "web_source", template_name) if not os.path.isfile(path): 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 _is_session_user_valid(): """校验 session 中的 user_id 是否在数据库中仍存在;不存在则清除 session 并返回 False""" uid = session.get('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() if not row: session.clear() return False return True except Exception: session.clear() return False def _get_current_admin_role(): """获取当前登录用户的管理角色:super_admin / admin / 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", (session['user_id'],) ) 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 login_required(f): @wraps(f) def decorated(*args, **kwargs): if not session.get('user_id') or not _is_session_user_valid(): if request.headers.get('X-Requested-With') == 'XMLHttpRequest': return jsonify({'success': False, 'error': '未登录'}), 401 return redirect(url_for('auth.login')) return f(*args, **kwargs) return decorated def admin_required(f): @wraps(f) def decorated(*args, **kwargs): if not session.get('user_id'): if request.headers.get('X-Requested-With') == 'XMLHttpRequest': return jsonify({'success': False, 'error': '未登录'}), 401 return redirect(url_for('auth.login')) try: conn = get_db() with conn.cursor() as cur: cur.execute("SELECT is_admin, role FROM users WHERE id = %s", (session['user_id'],)) row = cur.fetchone() conn.close() if not row or not row.get('is_admin'): if request.headers.get('X-Requested-With') == 'XMLHttpRequest': return jsonify({'success': False, 'error': '需要管理员权限'}), 403 return redirect(url_for('main.home')) except Exception as e: if request.headers.get('X-Requested-With') == 'XMLHttpRequest': 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