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 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,
                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,
                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_type TEXT NOT NULL,
                screen_id INTEGER,
                title 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_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);
            """
        )
        seed_default_tags(conn)


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),
        )