zsglpt/db/schema.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from __future__ import annotations

import sqlite3

from db.utils import get_cst_now_str


def ensure_schema(conn) -> None:
    """创建当前版本所需的所有表与索引（新库可直接得到完整 schema）。"""
    cursor = conn.cursor()

    # 管理员表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS admins (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """
    )

    # 用户表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            email TEXT,
            email_verified INTEGER DEFAULT 0,
            email_notify_enabled INTEGER DEFAULT 1,
            status TEXT DEFAULT 'approved',
            vip_expire_time TIMESTAMP,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            approved_at TIMESTAMP
        )
        """
    )

    # 登录设备指纹表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS login_fingerprints (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            user_agent TEXT NOT NULL,
            first_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            last_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            last_ip TEXT DEFAULT '',
            UNIQUE (user_id, user_agent),
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
        )
        """
    )

    # 登录IP记录表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS login_ips (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            ip TEXT NOT NULL,
            first_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            last_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            UNIQUE (user_id, ip),
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
        )
        """
    )

    # ==================== 安全防护：威胁检测相关表 ====================

    # 威胁事件日志表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS threat_events (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            threat_type TEXT NOT NULL,
            score INTEGER NOT NULL DEFAULT 0,
            rule TEXT,
            field_name TEXT,
            matched TEXT,
            value_preview TEXT,
            ip TEXT,
            user_id INTEGER,
            request_method TEXT,
            request_path TEXT,
            user_agent TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
        )
        """
    )

    # IP风险评分表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS ip_risk_scores (
            ip TEXT PRIMARY KEY,
            risk_score INTEGER NOT NULL DEFAULT 0,
            last_seen TIMESTAMP,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """
    )

    # 用户风险评分表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS user_risk_scores (
            user_id INTEGER PRIMARY KEY,
            risk_score INTEGER NOT NULL DEFAULT 0,
            last_seen TIMESTAMP,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
        )
        """
    )

    # IP黑名单表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS ip_blacklist (
            ip TEXT PRIMARY KEY,
            reason TEXT,
            is_active INTEGER DEFAULT 1,
            added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            expires_at TIMESTAMP
        )
        """
    )

    # 用户黑名单表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS user_blacklist (
            user_id INTEGER PRIMARY KEY,
            reason TEXT,
            is_active INTEGER DEFAULT 1,
            added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            expires_at TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
        )
        """
    )

    # 威胁特征库表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS threat_signatures (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT UNIQUE NOT NULL,
            threat_type TEXT NOT NULL,
            pattern TEXT NOT NULL,
            pattern_type TEXT DEFAULT 'regex',
            score INTEGER DEFAULT 0,
            is_active INTEGER DEFAULT 1,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """
    )

    # 账号表(关联用户)
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS accounts (
            id TEXT PRIMARY KEY,
            user_id INTEGER NOT NULL,
            username TEXT NOT NULL,
            password TEXT NOT NULL,
            remember INTEGER DEFAULT 1,
            remark TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
        )
        """
    )

    # VIP配置表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS vip_config (
            id INTEGER PRIMARY KEY CHECK (id = 1),
            default_vip_days INTEGER DEFAULT 0,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """
    )

    # 系统配置表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS system_config (
            id INTEGER PRIMARY KEY CHECK (id = 1),
            max_concurrent_global INTEGER DEFAULT 2,
            schedule_enabled INTEGER DEFAULT 0,
            schedule_time TEXT DEFAULT '02:00',
            schedule_browse_type TEXT DEFAULT '应读',
            proxy_enabled INTEGER DEFAULT 0,
            proxy_api_url TEXT DEFAULT '',
            proxy_expire_minutes INTEGER DEFAULT 3,
            max_screenshot_concurrent INTEGER DEFAULT 3,
            max_concurrent_per_account INTEGER DEFAULT 1,
            schedule_weekdays TEXT DEFAULT '1,2,3,4,5,6,7',
            enable_screenshot INTEGER DEFAULT 1,
            auto_approve_enabled INTEGER DEFAULT 0,
            auto_approve_hourly_limit INTEGER DEFAULT 10,
            auto_approve_vip_days INTEGER DEFAULT 7,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """
    )

    # 任务日志表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS task_logs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            account_id TEXT NOT NULL,
            username TEXT NOT NULL,
            browse_type TEXT NOT NULL,
            status TEXT NOT NULL,
            total_items INTEGER DEFAULT 0,
            total_attachments INTEGER DEFAULT 0,
            error_message TEXT,
            duration INTEGER,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            source TEXT DEFAULT 'manual',
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
        )
        """
    )

    # 数据库版本表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS db_version (
            id INTEGER PRIMARY KEY CHECK (id = 1),
            version INTEGER NOT NULL,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """
    )

    # Bug反馈表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS bug_feedbacks (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            username TEXT NOT NULL,
            title TEXT NOT NULL,
            description TEXT NOT NULL,
            contact TEXT,
            status TEXT DEFAULT 'pending',
            admin_reply TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            replied_at TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
        )
        """
    )

    # 公告表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS announcements (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            title TEXT NOT NULL,
            content TEXT NOT NULL,
            image_url TEXT,
            is_active INTEGER DEFAULT 1,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """
    )

    # 公告永久关闭记录表（用户维度）
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS announcement_dismissals (
            user_id INTEGER NOT NULL,
            announcement_id INTEGER NOT NULL,
            dismissed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            PRIMARY KEY (user_id, announcement_id),
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE,
            FOREIGN KEY (announcement_id) REFERENCES announcements (id) ON DELETE CASCADE
        )
        """
    )

    # 用户定时任务表
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS user_schedules (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            name TEXT DEFAULT '我的定时任务',
            enabled INTEGER DEFAULT 0,
            schedule_time TEXT NOT NULL DEFAULT '08:00',
            weekdays TEXT NOT NULL DEFAULT '1,2,3,4,5',
            browse_type TEXT NOT NULL DEFAULT '应读',
            enable_screenshot INTEGER DEFAULT 1,
            random_delay INTEGER DEFAULT 0,
            account_ids TEXT,
            last_run_at TIMESTAMP,
            next_run_at TIMESTAMP,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
        )
        """
    )

    # 定时任务执行日志表（历史上在迁移中创建；这里补齐，避免新库缺表）
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS schedule_execution_logs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            schedule_id INTEGER NOT NULL,
            user_id INTEGER NOT NULL,
            schedule_name TEXT,
            execute_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            total_accounts INTEGER DEFAULT 0,
            success_accounts INTEGER DEFAULT 0,
            failed_accounts INTEGER DEFAULT 0,
            total_items INTEGER DEFAULT 0,
            total_attachments INTEGER DEFAULT 0,
            total_screenshots INTEGER DEFAULT 0,
            duration_seconds INTEGER DEFAULT 0,
            status TEXT DEFAULT 'running',
            error_message TEXT,
            FOREIGN KEY (schedule_id) REFERENCES user_schedules (id) ON DELETE CASCADE,
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
        )
        """
    )

    # ========== 创建索引 ==========
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_users_username ON users(username)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_users_status ON users(status)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_users_vip_expire ON users(vip_expire_time)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_login_fingerprints_user ON login_fingerprints(user_id)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_login_ips_user ON login_ips(user_id)")

    cursor.execute("CREATE INDEX IF NOT EXISTS idx_threat_events_created_at ON threat_events(created_at)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_threat_events_ip ON threat_events(ip)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_threat_events_user_id ON threat_events(user_id)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_threat_events_type ON threat_events(threat_type)")

    cursor.execute("CREATE INDEX IF NOT EXISTS idx_ip_risk_scores_score ON ip_risk_scores(risk_score)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_ip_risk_scores_updated_at ON ip_risk_scores(updated_at)")

    cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_risk_scores_score ON user_risk_scores(risk_score)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_risk_scores_updated_at ON user_risk_scores(updated_at)")

    cursor.execute("CREATE INDEX IF NOT EXISTS idx_ip_blacklist_active ON ip_blacklist(is_active)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_ip_blacklist_expires ON ip_blacklist(expires_at)")

    cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_blacklist_active ON user_blacklist(is_active)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_blacklist_expires ON user_blacklist(expires_at)")

    cursor.execute("CREATE INDEX IF NOT EXISTS idx_threat_signatures_type ON threat_signatures(threat_type)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_threat_signatures_active ON threat_signatures(is_active)")

    cursor.execute("CREATE INDEX IF NOT EXISTS idx_accounts_user_id ON accounts(user_id)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_accounts_username ON accounts(username)")

    cursor.execute("CREATE INDEX IF NOT EXISTS idx_task_logs_user_id ON task_logs(user_id)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_task_logs_status ON task_logs(status)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_task_logs_created_at ON task_logs(created_at)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_task_logs_user_date ON task_logs(user_id, created_at)")

    cursor.execute("CREATE INDEX IF NOT EXISTS idx_bug_feedbacks_user_id ON bug_feedbacks(user_id)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_bug_feedbacks_status ON bug_feedbacks(status)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_bug_feedbacks_created_at ON bug_feedbacks(created_at)")

    cursor.execute("CREATE INDEX IF NOT EXISTS idx_announcements_active ON announcements(is_active)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_announcements_created_at ON announcements(created_at)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_announcement_dismissals_user ON announcement_dismissals(user_id)")

    cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_schedules_user_id ON user_schedules(user_id)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_schedules_enabled ON user_schedules(enabled)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_schedules_next_run ON user_schedules(next_run_at)")
    # 复合索引优化
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_schedules_user_enabled ON user_schedules(user_id, enabled)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_schedule_execution_logs_schedule_id ON schedule_execution_logs(schedule_id)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_schedule_execution_logs_user_id ON schedule_execution_logs(user_id)")
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_schedule_execution_logs_status ON schedule_execution_logs(status)")

    # 初始化VIP配置（幂等）
    try:
        cursor.execute("INSERT INTO vip_config (id, default_vip_days) VALUES (1, 0)")
    except sqlite3.IntegrityError:
        pass

    # 初始化系统配置（幂等）
    try:
        cursor.execute(
            """
            INSERT INTO system_config (
                id, max_concurrent_global, max_concurrent_per_account, max_screenshot_concurrent,
                schedule_enabled, schedule_time, schedule_browse_type, schedule_weekdays,
                proxy_enabled, proxy_api_url, proxy_expire_minutes, enable_screenshot,
                auto_approve_enabled, auto_approve_hourly_limit, auto_approve_vip_days
            ) VALUES (1, 2, 1, 3, 0, '02:00', '应读', '1,2,3,4,5,6,7', 0, '', 3, 1, 0, 10, 7)
            """
        )
    except sqlite3.IntegrityError:
        pass

    # 确保 db_version 记录存在（默认 0，由迁移统一更新）
    cursor.execute("INSERT OR IGNORE INTO db_version (id, version, updated_at) VALUES (1, 0, ?)", (get_cst_now_str(),))

    conn.commit()