win_monitor/backend/app/database.py

from __future__ import annotations

import sqlite3
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator


BASE_DIR = Path(__file__).resolve().parent.parent
DATA_DIR = BASE_DIR / "data"
DB_PATH = DATA_DIR / "win_monitor.db"


def dict_factory(cursor: sqlite3.Cursor, row: sqlite3.Row) -> dict:
    fields = [column[0] for column in cursor.description]
    return {key: row[index] for index, key in enumerate(fields)}


@contextmanager
def get_db() -> Iterator[sqlite3.Connection]:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = dict_factory
    conn.execute("PRAGMA foreign_keys = ON")
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def init_db() -> None:
    with get_db() as conn:
        conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS scan_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                started_at TEXT NOT NULL,
                finished_at TEXT,
                status TEXT NOT NULL,
                services_found INTEGER NOT NULL DEFAULT 0,
                processes_found INTEGER NOT NULL DEFAULT 0,
                new_services INTEGER NOT NULL DEFAULT 0,
                new_processes INTEGER NOT NULL DEFAULT 0,
                error_message TEXT
            );

            CREATE TABLE IF NOT EXISTS windows_services (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                identity_key TEXT NOT NULL UNIQUE,
                name TEXT NOT NULL,
                display_name TEXT,
                status TEXT,
                start_type TEXT,
                username TEXT,
                binary_path TEXT,
                description TEXT,
                is_present_now INTEGER NOT NULL DEFAULT 1,
                first_seen_at TEXT NOT NULL,
                last_seen_at TEXT NOT NULL,
                confirm_status TEXT NOT NULL DEFAULT 'PENDING',
                user_note TEXT,
                ai_description TEXT,
                ai_reason TEXT,
                ai_suggestion TEXT,
                risk_level TEXT,
                updated_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS windows_processes (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                identity_key TEXT NOT NULL UNIQUE,
                name TEXT NOT NULL,
                exe_path TEXT,
                cmdline TEXT,
                username TEXT,
                status TEXT,
                cwd TEXT,
                last_pid INTEGER,
                parent_pid INTEGER,
                create_time TEXT,
                is_present_now INTEGER NOT NULL DEFAULT 1,
                first_seen_at TEXT NOT NULL,
                last_seen_at TEXT NOT NULL,
                confirm_status TEXT NOT NULL DEFAULT 'PENDING',
                user_note TEXT,
                ai_description TEXT,
                ai_reason TEXT,
                ai_suggestion TEXT,
                risk_level TEXT,
                updated_at TEXT NOT NULL
            );

            CREATE INDEX IF NOT EXISTS idx_services_confirm_status
                ON windows_services(confirm_status);
            CREATE INDEX IF NOT EXISTS idx_processes_confirm_status
                ON windows_processes(confirm_status);
            CREATE INDEX IF NOT EXISTS idx_scan_records_started_at
                ON scan_records(started_at DESC);

            CREATE TABLE IF NOT EXISTS tags (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL UNIQUE,
                description TEXT,
                is_controllable INTEGER NOT NULL DEFAULT 1,
                is_builtin INTEGER NOT NULL DEFAULT 0,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS item_tags (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                item_type TEXT NOT NULL CHECK(item_type IN ('service', 'process')),
                item_id INTEGER NOT NULL,
                tag_id INTEGER NOT NULL,
                created_at TEXT NOT NULL,
                UNIQUE(item_type, item_id, tag_id),
                FOREIGN KEY(tag_id) REFERENCES tags(id) ON DELETE CASCADE
            );

            CREATE INDEX IF NOT EXISTS idx_item_tags_item
                ON item_tags(item_type, item_id);

            CREATE TABLE IF NOT EXISTS ai_providers (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL UNIQUE,
                provider_type TEXT NOT NULL CHECK(provider_type IN ('OPENAI', 'OPENAI_COMPATIBLE', 'GOOGLE_GEMINI')),
                base_url TEXT,
                api_key TEXT,
                enabled INTEGER NOT NULL DEFAULT 1,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS ai_models (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                provider_id INTEGER NOT NULL,
                name TEXT NOT NULL,
                display_name TEXT,
                is_default INTEGER NOT NULL DEFAULT 0,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                UNIQUE(provider_id, name),
                FOREIGN KEY(provider_id) REFERENCES ai_providers(id) ON DELETE CASCADE
            );

            CREATE INDEX IF NOT EXISTS idx_ai_models_provider
                ON ai_models(provider_id);

            CREATE TABLE IF NOT EXISTS app_settings (
                key TEXT PRIMARY KEY,
                value TEXT,
                description TEXT,
                updated_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS automation_screens (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                interface_name TEXT NOT NULL,
                description TEXT,
                image_path TEXT NOT NULL,
                width INTEGER NOT NULL,
                height INTEGER NOT NULL,
                is_windows_desktop INTEGER NOT NULL DEFAULT 0,
                is_browser_webpage INTEGER NOT NULL DEFAULT 0,
                raw_ai_json TEXT,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS automation_screen_elements (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                screen_id INTEGER NOT NULL,
                element_index INTEGER NOT NULL,
                name TEXT NOT NULL,
                x_percent REAL NOT NULL,
                y_percent REAL NOT NULL,
                x INTEGER NOT NULL,
                y INTEGER NOT NULL,
                approximate_location TEXT,
                is_located INTEGER NOT NULL DEFAULT 0,
                raw_json TEXT,
                created_at TEXT NOT NULL,
                FOREIGN KEY(screen_id) REFERENCES automation_screens(id) ON DELETE CASCADE
            );

            CREATE INDEX IF NOT EXISTS idx_automation_screen_elements_screen
                ON automation_screen_elements(screen_id, element_index);

            CREATE TABLE IF NOT EXISTS automation_workflows (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                workflow_key TEXT UNIQUE,
                name TEXT NOT NULL,
                description TEXT,
                raw_json TEXT,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS automation_workflow_nodes (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                workflow_id INTEGER NOT NULL,
                node_index INTEGER NOT NULL,
                node_key TEXT,
                node_type TEXT NOT NULL,
                screen_id INTEGER,
                title TEXT,
                position_x REAL NOT NULL DEFAULT 80,
                position_y REAL NOT NULL DEFAULT 80,
                next_node_keys TEXT,
                config_json TEXT NOT NULL,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                FOREIGN KEY(workflow_id) REFERENCES automation_workflows(id) ON DELETE CASCADE,
                FOREIGN KEY(screen_id) REFERENCES automation_screens(id) ON DELETE SET NULL
            );

            CREATE INDEX IF NOT EXISTS idx_automation_workflow_nodes_workflow
                ON automation_workflow_nodes(workflow_id, node_index);

            CREATE TABLE IF NOT EXISTS automation_workflow_tasks (
                id TEXT PRIMARY KEY,
                workflow_id INTEGER,
                workflow_key TEXT NOT NULL,
                workflow_name TEXT NOT NULL,
                status TEXT NOT NULL,
                request_json TEXT NOT NULL,
                workflow_snapshot_json TEXT NOT NULL,
                result_json TEXT,
                return_data_json TEXT,
                error_message TEXT,
                created_at TEXT NOT NULL,
                started_at TEXT,
                finished_at TEXT,
                FOREIGN KEY(workflow_id) REFERENCES automation_workflows(id) ON DELETE SET NULL
            );

            CREATE INDEX IF NOT EXISTS idx_automation_workflow_tasks_status_created
                ON automation_workflow_tasks(status, created_at);
            CREATE INDEX IF NOT EXISTS idx_automation_workflow_tasks_created
                ON automation_workflow_tasks(created_at DESC);

            CREATE TABLE IF NOT EXISTS automation_workflow_runtime (
                id INTEGER PRIMARY KEY CHECK(id = 1),
                active_task_id TEXT,
                updated_at TEXT,
                FOREIGN KEY(active_task_id) REFERENCES automation_workflow_tasks(id) ON DELETE SET NULL
            );

            INSERT OR IGNORE INTO automation_workflow_runtime (id, active_task_id, updated_at)
            VALUES (1, NULL, NULL);

            CREATE TABLE IF NOT EXISTS automation_errors (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                workflow_id INTEGER,
                node_id INTEGER,
                screen_id INTEGER,
                action_type TEXT,
                message TEXT NOT NULL,
                similarity REAL,
                expected_image_path TEXT,
                actual_image_path TEXT,
                compare_result_json TEXT,
                created_at TEXT NOT NULL,
                FOREIGN KEY(workflow_id) REFERENCES automation_workflows(id) ON DELETE SET NULL,
                FOREIGN KEY(node_id) REFERENCES automation_workflow_nodes(id) ON DELETE SET NULL,
                FOREIGN KEY(screen_id) REFERENCES automation_screens(id) ON DELETE SET NULL
            );

            CREATE INDEX IF NOT EXISTS idx_automation_errors_created
                ON automation_errors(created_at DESC);
            """
        )
        ensure_column(conn, "automation_workflows", "workflow_key", "TEXT")
        conn.execute(
            """
            CREATE UNIQUE INDEX IF NOT EXISTS idx_automation_workflows_key
                ON automation_workflows(workflow_key)
                WHERE workflow_key IS NOT NULL AND workflow_key != ''
            """
        )
        ensure_column(conn, "automation_workflow_nodes", "node_key", "TEXT")
        ensure_column(conn, "automation_workflow_nodes", "position_x", "REAL NOT NULL DEFAULT 80")
        ensure_column(conn, "automation_workflow_nodes", "position_y", "REAL NOT NULL DEFAULT 80")
        ensure_column(conn, "automation_workflow_nodes", "next_node_keys", "TEXT")
        ensure_column(conn, "automation_screen_elements", "approximate_location", "TEXT")
        ensure_column(conn, "automation_screen_elements", "is_located", "INTEGER NOT NULL DEFAULT 0")
        seed_default_tags(conn)
        seed_default_settings(conn)


def ensure_column(conn: sqlite3.Connection, table: str, column: str, definition: str) -> None:
    existing = {row["name"] for row in conn.execute(f"PRAGMA table_info({table})").fetchall()}
    if column not in existing:
        conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {definition}")


def seed_default_tags(conn: sqlite3.Connection) -> None:
    rows = [
        (
            "windows系统",
            "Windows 系统原生服务或进程。默认不可控制，避免误停止系统关键组件。",
            0,
            1,
        ),
        (
            "本系统相关",
            "本项目程序自身相关服务或进程。默认不可控制，避免通过监控系统停止自身。",
            0,
            1,
        ),
    ]
    now = __import__("datetime").datetime.now().astimezone().isoformat(timespec="seconds")
    for name, description, is_controllable, is_builtin in rows:
        conn.execute(
            """
            INSERT INTO tags (name, description, is_controllable, is_builtin, created_at, updated_at)
            VALUES (?, ?, ?, ?, ?, ?)
            ON CONFLICT(name) DO UPDATE SET
                description = COALESCE(tags.description, excluded.description),
                is_builtin = 1,
                updated_at = excluded.updated_at
            """,
            (name, description, is_controllable, is_builtin, now, now),
        )


def seed_default_settings(conn: sqlite3.Connection) -> None:
    settings = [
        ("default_ai_provider_id", "", "全局默认 AI 服务商 ID"),
        ("default_ai_model_id", "", "全局默认 AI 模型 ID"),
        ("default_ai_temperature", "0.1", "全局默认 AI 温度参数"),
        ("automation_file_root", "automation", "自动化文件保存根路径"),
        ("automation_screen_path", "automation/screens", "已识别界面截图保存路径"),
        ("automation_error_path", "automation/errors", "自动化错误截图保存路径"),
        ("automation_runtime_path", "automation/runtime", "自动化临时截图保存路径"),
        ("automation_auto_screenshot_enabled", "0", "自动化操作页面是否默认自动截屏"),
        ("automation_auto_screenshot_interval", "30", "自动化操作页面默认自动截屏间隔秒数"),
        ("automation_remote_token", "", "远程执行工作流和查询任务状态的 Token，设置后需通过 X-Automation-Token 或 Bearer Token 传入"),
    ]
    now = __import__("datetime").datetime.now().astimezone().isoformat(timespec="seconds")
    for key, value, description in settings:
        conn.execute(
            """
            INSERT INTO app_settings (key, value, description, updated_at)
            VALUES (?, ?, ?, ?)
            ON CONFLICT(key) DO UPDATE SET
                description = excluded.description
            """,
            (key, value, description, now),
        )