crawler-plugin/app/app_common.py

"""
公共模块：数据库连接、初始化、会话校验、装饰器、模板渲染
供各蓝图复用
"""
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