win_monitor/backend/app/main.py

from __future__ import annotations

import json
import sqlite3
from typing import Any

from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware

from .control import (
    CONFIRMED_CONTROL_STATUSES,
    restart_service,
    start_process,
    start_service,
    stop_process,
    stop_service,
)
from .database import get_db, init_db
from .scanner import now_iso, run_full_scan
from .sensors import collect_sensors
from .schemas import AiImportRequest, BatchStatusUpdate, PromptRequest, StatusUpdate, TagAssignRequest, TagCreate, TagUpdate
from .smart import collect_all_smart, get_device_smart, scan_devices


AI_PROMPT_TEMPLATE = """请作为资深的 Windows 系统安全专家，帮我分析下面这些 Windows 服务和进程是否可信，并严格按照 JSON 数组格式输出结果。

输出要求：
1. 必须且只能输出纯 JSON 数组，不要输出任何额外的解释、问候语，也不要使用 Markdown 代码块（如 ```json）包裹。
2. 每个对象必须包含以下 7 个字段：type、name、description、judgement、risk_level、reason、suggestion。
3. type 只能是 "service" 或 "process"。
4. description 请简要说明该服务或进程的官方用途或常规功能（如果是未知/恶意程序，请描述其伪装意图或表现）。
5. judgement 只能是 "TRUSTED"、"SUSPICIOUS"、"NEED_MORE_INFO"。
6. risk_level 只能是 "LOW"、"MEDIUM"、"HIGH"。
7. 如果提供的信息不足以做出判断，请将 judgement 设为 "NEED_MORE_INFO"。
8. 请结合每个对象的 tags 字段进行判断。已有标签是人工上下文，不代表最终结论，但如果标签显示为“windows系统”或“本系统相关”，请在 reason 或 suggestion 中体现这一点。
9. 如果你认为某个对象适合系统已有标签，可以在 suggestion 中建议使用对应标签名称；不要创造不存在的新标签。

JSON 格式示例：

[
  {
    "type": "service",
    "name": "WinDefend",
    "description": "Microsoft Defender 防病毒核心服务，负责保护系统免受恶意软件和间谍软件的威胁。",
    "judgement": "TRUSTED",
    "risk_level": "LOW",
    "reason": "这是 Microsoft 官方的安全组件，路径和名称符合系统原生服务的标准特征。",
    "suggestion": "可标记为可信，建议保持运行。"
  },
  {
    "type": "process",
    "name": "unknown.exe",
    "description": "未知用途的执行文件，无明确的官方功能说明。",
    "judgement": "SUSPICIOUS",
    "risk_level": "HIGH",
    "reason": "进程位于用户 AppData 临时目录，启动命令行异常，且缺少有效的官方数字签名。",
    "suggestion": "建议立即隔离，检查文件的 SHA256 散列值及外部网络连接记录，不要直接运行或信任。"
  }
]

下面是待分析数据：

{pending_items_json}

系统中已有标签信息：

{tags_json}
"""


app = FastAPI(title="Windows Monitor API", version="1.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("startup")
def startup() -> None:
    init_db()


def build_where(
    keyword: str | None,
    confirm_status: str | None,
    present: bool | None,
    fields: list[str],
) -> tuple[str, list[Any]]:
    clauses: list[str] = []
    params: list[Any] = []
    if keyword:
        like = f"%{keyword}%"
        clauses.append("(" + " OR ".join(f"{field} LIKE ?" for field in fields) + ")")
        params.extend([like] * len(fields))
    if confirm_status:
        clauses.append("confirm_status = ?")
        params.append(confirm_status)
    if present is not None:
        clauses.append("is_present_now = ?")
        params.append(1 if present else 0)
    return ("WHERE " + " AND ".join(clauses)) if clauses else "", params


def list_items(
    table: str,
    keyword: str | None,
    confirm_status: str | None,
    present: bool | None,
    page: int,
    page_size: int,
    fields: list[str],
) -> dict[str, Any]:
    where_sql, params = build_where(keyword, confirm_status, present, fields)
    offset = (page - 1) * page_size
    with get_db() as conn:
        total = conn.execute(f"SELECT COUNT(*) AS total FROM {table} {where_sql}", params).fetchone()["total"]
        rows = conn.execute(
            f"SELECT * FROM {table} {where_sql} ORDER BY is_present_now DESC, last_seen_at DESC LIMIT ? OFFSET ?",
            [*params, page_size, offset],
        ).fetchall()
        rows = attach_item_metadata(conn, table_to_item_type(table), rows)
    return {"items": rows, "total": total, "page": page, "page_size": page_size}


def get_item(table: str, item_id: int) -> dict[str, Any]:
    with get_db() as conn:
        item = conn.execute(f"SELECT * FROM {table} WHERE id = ?", (item_id,)).fetchone()
        if item:
            item = attach_item_metadata(conn, table_to_item_type(table), [item])[0]
    if not item:
        raise HTTPException(status_code=404, detail="Item not found")
    return item


def table_to_item_type(table: str) -> str:
    if table == "windows_services":
        return "service"
    if table == "windows_processes":
        return "process"
    raise ValueError(f"Unsupported table: {table}")


def bool_tag(row: dict[str, Any]) -> dict[str, Any]:
    item = dict(row)
    item["is_controllable"] = bool(item["is_controllable"])
    item["is_builtin"] = bool(item["is_builtin"])
    return item


def all_tags(conn) -> list[dict[str, Any]]:
    return [
        bool_tag(row)
        for row in conn.execute("SELECT * FROM tags ORDER BY is_builtin DESC, name ASC").fetchall()
    ]


def tags_for_items(conn, item_type: str, item_ids: list[int]) -> dict[int, list[dict[str, Any]]]:
    if not item_ids:
        return {}
    placeholders = ",".join("?" for _ in item_ids)
    rows = conn.execute(
        f"""
        SELECT it.item_id, t.*
        FROM item_tags it
        JOIN tags t ON t.id = it.tag_id
        WHERE it.item_type = ? AND it.item_id IN ({placeholders})
        ORDER BY t.name ASC
        """,
        [item_type, *item_ids],
    ).fetchall()
    result = {item_id: [] for item_id in item_ids}
    for row in rows:
        item_id = row["item_id"]
        tag = {key: value for key, value in row.items() if key != "item_id"}
        result.setdefault(item_id, []).append(bool_tag(tag))
    return result


def can_control_item(item_type: str, row: dict[str, Any], tags: list[dict[str, Any]]) -> bool:
    if row.get("confirm_status") not in CONFIRMED_CONTROL_STATUSES:
        return False
    if item_type == "process":
        protected_names = {"system idle process", "system", "registry"}
        if row.get("last_pid") in (0, 4) or (row.get("name") or "").lower() in protected_names:
            return False
    return all(tag.get("is_controllable", True) for tag in tags)


def attach_item_metadata(conn, item_type: str, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    tag_map = tags_for_items(conn, item_type, [row["id"] for row in rows])
    enriched = []
    for row in rows:
        item = dict(row)
        item["is_present_now"] = bool(item.get("is_present_now"))
        item["tags"] = tag_map.get(row["id"], [])
        item["can_control"] = can_control_item(item_type, item, item["tags"])
        enriched.append(item)
    return enriched


def update_one(table: str, item_id: int, payload: StatusUpdate) -> dict[str, Any]:
    get_item(table, item_id)
    with get_db() as conn:
        conn.execute(
            f"UPDATE {table} SET confirm_status = ?, user_note = ?, updated_at = ? WHERE id = ?",
            (payload.confirm_status, payload.user_note, now_iso(), item_id),
        )
    return get_item(table, item_id)


def update_batch(table: str, payload: BatchStatusUpdate) -> dict[str, Any]:
    if not payload.ids:
        return {"updated": 0}
    placeholders = ",".join("?" for _ in payload.ids)
    with get_db() as conn:
        cursor = conn.execute(
            f"""
            UPDATE {table}
            SET confirm_status = ?, user_note = COALESCE(?, user_note), updated_at = ?
            WHERE id IN ({placeholders})
            """,
            [payload.confirm_status, payload.user_note, now_iso(), *payload.ids],
        )
    return {"updated": cursor.rowcount}


def rows_for_prompt(table: str, item_type: str, payload: PromptRequest) -> list[dict[str, Any]]:
    with get_db() as conn:
        if payload.scope == "selected" and payload.ids:
            placeholders = ",".join("?" for _ in payload.ids)
            rows = conn.execute(f"SELECT * FROM {table} WHERE id IN ({placeholders})", payload.ids).fetchall()
        else:
            rows = conn.execute(f"SELECT * FROM {table} WHERE confirm_status = 'PENDING'").fetchall()
        rows = attach_item_metadata(conn, item_type, rows)
    return [normalize_prompt_row(item_type, row) for row in rows]


def normalize_prompt_row(item_type: str, row: dict[str, Any]) -> dict[str, Any]:
    if item_type == "service":
        return {
            "type": "service",
            "id": row["id"],
            "name": row["name"],
            "display_name": row["display_name"],
            "status": row["status"],
            "start_type": row["start_type"],
            "username": row["username"],
            "binary_path": row["binary_path"],
            "description": row["description"],
            "is_present_now": bool(row["is_present_now"]),
            "tags": [
                {
                    "name": tag["name"],
                    "description": tag["description"],
                    "is_controllable": tag["is_controllable"],
                }
                for tag in row.get("tags", [])
            ],
        }
    return {
        "type": "process",
        "id": row["id"],
        "name": row["name"],
        "exe_path": row["exe_path"],
        "cmdline": row["cmdline"],
        "username": row["username"],
        "status": row["status"],
        "last_pid": row["last_pid"],
        "parent_pid": row["parent_pid"],
        "is_present_now": bool(row["is_present_now"]),
        "tags": [
            {
                "name": tag["name"],
                "description": tag["description"],
                "is_controllable": tag["is_controllable"],
            }
            for tag in row.get("tags", [])
        ],
    }


def markdown_table(rows: list[dict[str, Any]]) -> str:
    headers = ["type", "id", "name", "status", "tags", "path_or_command", "user", "present"]
    lines = ["| " + " | ".join(headers) + " |", "| " + " | ".join(["---"] * len(headers)) + " |"]
    for row in rows:
        path_or_command = row.get("binary_path") or row.get("exe_path") or row.get("cmdline") or ""
        values = [
            str(row.get("type", "")),
            str(row.get("id", "")),
            str(row.get("name", "")),
            str(row.get("status", "")),
            ", ".join(tag.get("name", "") for tag in row.get("tags", [])).replace("|", "\\|"),
            str(path_or_command).replace("|", "\\|"),
            str(row.get("username", "")).replace("|", "\\|"),
            "yes" if row.get("is_present_now") else "no",
        ]
        lines.append("| " + " | ".join(values) + " |")
    return "\n".join(lines)


def prompt_response(rows: list[dict[str, Any]]) -> dict[str, Any]:
    pending_json = json.dumps(rows, ensure_ascii=False, indent=2)
    table = markdown_table(rows)
    with get_db() as conn:
        tags_json = json.dumps(all_tags(conn), ensure_ascii=False, indent=2)
    prompt_text = AI_PROMPT_TEMPLATE.replace("{pending_items_json}", pending_json).replace("{tags_json}", tags_json)
    return {"prompt_text": prompt_text, "markdown_table": table, "items": rows}


def set_item_tags(item_type: str, table: str, item_id: int, payload: TagAssignRequest) -> dict[str, Any]:
    get_item(table, item_id)
    unique_ids = sorted(set(payload.tag_ids))
    now = now_iso()
    with get_db() as conn:
        if unique_ids:
            placeholders = ",".join("?" for _ in unique_ids)
            found = conn.execute(f"SELECT id FROM tags WHERE id IN ({placeholders})", unique_ids).fetchall()
            if len(found) != len(unique_ids):
                raise HTTPException(status_code=400, detail="One or more tag ids do not exist")
        conn.execute("DELETE FROM item_tags WHERE item_type = ? AND item_id = ?", (item_type, item_id))
        for tag_id in unique_ids:
            conn.execute(
                "INSERT INTO item_tags (item_type, item_id, tag_id, created_at) VALUES (?, ?, ?, ?)",
                (item_type, item_id, tag_id, now),
            )
    return get_item(table, item_id)


def ensure_control_allowed(table: str, item_id: int) -> dict[str, Any]:
    item = get_item(table, item_id)
    if not item.get("can_control"):
        raise HTTPException(status_code=403, detail="This item is not controllable because it is unconfirmed or has a non-controllable tag")
    return item


def import_ai_results(table: str, item_type: str, payload: AiImportRequest) -> dict[str, Any]:
    updated = 0
    with get_db() as conn:
        for item in payload.items:
            if item.type != item_type:
                continue
            cursor = conn.execute(
                f"""
                UPDATE {table}
                SET confirm_status = ?, ai_description = ?, ai_reason = ?,
                    ai_suggestion = ?, risk_level = ?, updated_at = ?
                WHERE name = ?
                """,
                (
                    item.judgement,
                    item.description,
                    item.reason,
                    item.suggestion,
                    item.risk_level,
                    now_iso(),
                    item.name,
                ),
            )
            updated += cursor.rowcount
    return {"updated": updated}


@app.get("/api/dashboard")
def dashboard() -> dict[str, Any]:
    with get_db() as conn:
        latest_scan = conn.execute("SELECT * FROM scan_records ORDER BY started_at DESC LIMIT 1").fetchone()
        service_total = conn.execute("SELECT COUNT(*) AS total FROM windows_services").fetchone()["total"]
        process_total = conn.execute("SELECT COUNT(*) AS total FROM windows_processes").fetchone()["total"]
        pending_services = conn.execute(
            "SELECT COUNT(*) AS total FROM windows_services WHERE confirm_status = 'PENDING'"
        ).fetchone()["total"]
        pending_processes = conn.execute(
            "SELECT COUNT(*) AS total FROM windows_processes WHERE confirm_status = 'PENDING'"
        ).fetchone()["total"]
        missing_services = conn.execute(
            "SELECT COUNT(*) AS total FROM windows_services WHERE is_present_now = 0"
        ).fetchone()["total"]
        missing_processes = conn.execute(
            "SELECT COUNT(*) AS total FROM windows_processes WHERE is_present_now = 0"
        ).fetchone()["total"]
    return {
        "latest_scan": latest_scan,
        "service_total": service_total,
        "process_total": process_total,
        "pending_services": pending_services,
        "pending_processes": pending_processes,
        "missing_services": missing_services,
        "missing_processes": missing_processes,
    }


@app.get("/api/tags")
def tags() -> dict[str, Any]:
    with get_db() as conn:
        rows = all_tags(conn)
    return {"items": rows}


@app.post("/api/tags")
def tag_create(payload: TagCreate) -> dict[str, Any]:
    now = now_iso()
    try:
        with get_db() as conn:
            cursor = conn.execute(
                """
                INSERT INTO tags (name, description, is_controllable, is_builtin, created_at, updated_at)
                VALUES (?, ?, ?, 0, ?, ?)
                """,
                (payload.name.strip(), payload.description, 1 if payload.is_controllable else 0, now, now),
            )
            tag_id = cursor.lastrowid
            row = conn.execute("SELECT * FROM tags WHERE id = ?", (tag_id,)).fetchone()
    except sqlite3.IntegrityError as exc:
        raise HTTPException(status_code=409, detail="Tag name already exists") from exc
    return bool_tag(row)


@app.patch("/api/tags/{tag_id}")
def tag_update(tag_id: int, payload: TagUpdate) -> dict[str, Any]:
    now = now_iso()
    try:
        with get_db() as conn:
            existing = conn.execute("SELECT * FROM tags WHERE id = ?", (tag_id,)).fetchone()
            if not existing:
                raise HTTPException(status_code=404, detail="Tag not found")
            conn.execute(
                """
                UPDATE tags
                SET name = ?, description = ?, is_controllable = ?, updated_at = ?
                WHERE id = ?
                """,
                (payload.name.strip(), payload.description, 1 if payload.is_controllable else 0, now, tag_id),
            )
            row = conn.execute("SELECT * FROM tags WHERE id = ?", (tag_id,)).fetchone()
    except sqlite3.IntegrityError as exc:
        raise HTTPException(status_code=409, detail="Tag name already exists") from exc
    return bool_tag(row)


@app.delete("/api/tags/{tag_id}")
def tag_delete(tag_id: int) -> dict[str, Any]:
    with get_db() as conn:
        row = conn.execute("SELECT * FROM tags WHERE id = ?", (tag_id,)).fetchone()
        if not row:
            raise HTTPException(status_code=404, detail="Tag not found")
        if row["is_builtin"]:
            raise HTTPException(status_code=400, detail="Built-in tags cannot be deleted")
        cursor = conn.execute("DELETE FROM tags WHERE id = ?", (tag_id,))
    return {"deleted": cursor.rowcount}


@app.get("/api/sensors")
def sensors() -> dict[str, Any]:
    return collect_sensors()


@app.get("/api/smart/scan")
def smart_scan() -> dict[str, Any]:
    return scan_devices()


@app.get("/api/smart/devices")
def smart_devices(include_jmb39x: bool = True, jmb39x_slots: int = Query(default=8, ge=0, le=16)) -> dict[str, Any]:
    return collect_all_smart(include_jmb39x=include_jmb39x, jmb39x_slots=jmb39x_slots)


@app.get("/api/smart/device")
def smart_device(device: str, device_type: str | None = None) -> dict[str, Any]:
    return get_device_smart(device, device_type)


@app.post("/api/scans/run")
def run_scan() -> dict[str, Any]:
    return run_full_scan()


@app.get("/api/scans")
def scan_history(page: int = 1, page_size: int = 20) -> dict[str, Any]:
    offset = (page - 1) * page_size
    with get_db() as conn:
        total = conn.execute("SELECT COUNT(*) AS total FROM scan_records").fetchone()["total"]
        rows = conn.execute(
            "SELECT * FROM scan_records ORDER BY started_at DESC LIMIT ? OFFSET ?",
            (page_size, offset),
        ).fetchall()
    return {"items": rows, "total": total, "page": page, "page_size": page_size}


@app.get("/api/scans/{scan_id}")
def scan_detail(scan_id: int) -> dict[str, Any]:
    with get_db() as conn:
        scan = conn.execute("SELECT * FROM scan_records WHERE id = ?", (scan_id,)).fetchone()
    if not scan:
        raise HTTPException(status_code=404, detail="Scan not found")
    return scan


@app.get("/api/services")
def services(
    keyword: str | None = None,
    confirm_status: str | None = None,
    present: bool | None = None,
    page: int = Query(default=1, ge=1),
    page_size: int = Query(default=20, ge=1, le=200),
) -> dict[str, Any]:
    return list_items(
        "windows_services",
        keyword,
        confirm_status,
        present,
        page,
        page_size,
        ["name", "display_name", "binary_path", "description"],
    )


@app.patch("/api/services/batch")
def service_batch_update(payload: BatchStatusUpdate) -> dict[str, Any]:
    return update_batch("windows_services", payload)


@app.post("/api/services/import-ai")
def service_import_ai(payload: AiImportRequest) -> dict[str, Any]:
    return import_ai_results("windows_services", "service", payload)


@app.post("/api/services/ai-prompt")
def service_ai_prompt(payload: PromptRequest) -> dict[str, Any]:
    return prompt_response(rows_for_prompt("windows_services", "service", payload))


@app.put("/api/services/{service_id}/tags")
def service_tags_update(service_id: int, payload: TagAssignRequest) -> dict[str, Any]:
    return set_item_tags("service", "windows_services", service_id, payload)


@app.post("/api/services/{service_id}/start")
def service_start(service_id: int) -> dict[str, Any]:
    item = ensure_control_allowed("windows_services", service_id)
    if not item.get("is_present_now"):
        raise HTTPException(status_code=400, detail="This service was not present in the latest scan")
    result = start_service(item["name"])
    with get_db() as conn:
        conn.execute(
            "UPDATE windows_services SET status = ?, updated_at = ? WHERE id = ?",
            (result.get("status") or "running", now_iso(), service_id),
        )
    return result


@app.post("/api/services/{service_id}/stop")
def service_stop(service_id: int) -> dict[str, Any]:
    item = ensure_control_allowed("windows_services", service_id)
    if not item.get("is_present_now"):
        raise HTTPException(status_code=400, detail="This service was not present in the latest scan")
    result = stop_service(item["name"])
    with get_db() as conn:
        conn.execute(
            "UPDATE windows_services SET status = ?, updated_at = ? WHERE id = ?",
            (result.get("status") or "stopped", now_iso(), service_id),
        )
    return result


@app.post("/api/services/{service_id}/restart")
def service_restart(service_id: int) -> dict[str, Any]:
    item = ensure_control_allowed("windows_services", service_id)
    if not item.get("is_present_now"):
        raise HTTPException(status_code=400, detail="This service was not present in the latest scan")
    result = restart_service(item["name"])
    with get_db() as conn:
        conn.execute(
            "UPDATE windows_services SET status = ?, updated_at = ? WHERE id = ?",
            (result.get("status") or "running", now_iso(), service_id),
        )
    return result


@app.get("/api/services/{service_id}")
def service_detail(service_id: int) -> dict[str, Any]:
    return get_item("windows_services", service_id)


@app.patch("/api/services/{service_id}")
def service_update(service_id: int, payload: StatusUpdate) -> dict[str, Any]:
    return update_one("windows_services", service_id, payload)


@app.get("/api/processes")
def processes(
    keyword: str | None = None,
    confirm_status: str | None = None,
    present: bool | None = None,
    page: int = Query(default=1, ge=1),
    page_size: int = Query(default=20, ge=1, le=200),
) -> dict[str, Any]:
    return list_items(
        "windows_processes",
        keyword,
        confirm_status,
        present,
        page,
        page_size,
        ["name", "exe_path", "cmdline", "username"],
    )


@app.patch("/api/processes/batch")
def process_batch_update(payload: BatchStatusUpdate) -> dict[str, Any]:
    return update_batch("windows_processes", payload)


@app.post("/api/processes/import-ai")
def process_import_ai(payload: AiImportRequest) -> dict[str, Any]:
    return import_ai_results("windows_processes", "process", payload)


@app.post("/api/processes/ai-prompt")
def process_ai_prompt(payload: PromptRequest) -> dict[str, Any]:
    return prompt_response(rows_for_prompt("windows_processes", "process", payload))


@app.put("/api/processes/{process_id}/tags")
def process_tags_update(process_id: int, payload: TagAssignRequest) -> dict[str, Any]:
    return set_item_tags("process", "windows_processes", process_id, payload)


@app.post("/api/processes/{process_id}/start")
def process_start(process_id: int) -> dict[str, Any]:
    item = ensure_control_allowed("windows_processes", process_id)
    result = start_process(item)
    with get_db() as conn:
        conn.execute(
            "UPDATE windows_processes SET last_pid = ?, is_present_now = 1, status = ?, updated_at = ? WHERE id = ?",
            (result.get("pid"), "running", now_iso(), process_id),
        )
    return result


@app.post("/api/processes/{process_id}/stop")
def process_stop(process_id: int) -> dict[str, Any]:
    item = ensure_control_allowed("windows_processes", process_id)
    if not item.get("is_present_now"):
        raise HTTPException(status_code=400, detail="This process was not present in the latest scan")
    result = stop_process(item)
    with get_db() as conn:
        conn.execute(
            "UPDATE windows_processes SET is_present_now = 0, status = ?, updated_at = ? WHERE id = ?",
            ("stopped", now_iso(), process_id),
        )
    return result


@app.get("/api/processes/{process_id}")
def process_detail(process_id: int) -> dict[str, Any]:
    return get_item("windows_processes", process_id)


@app.patch("/api/processes/{process_id}")
def process_update(process_id: int, payload: StatusUpdate) -> dict[str, Any]:
    return update_one("windows_processes", process_id, payload)