From 7f8dab4ec92f183afdda1d0fc8a517c72fa37852 Mon Sep 17 00:00:00 2001 From: Ed1s0nZ Date: Sat, 10 Oct 2026 20:45:39 +0800 Subject: [PATCH] feat: reconcile SQLite tables and columns at startup --- .../startup-schema-reconcile/design.md | 6 +- .../implementation-plan.md | 2 +- .../implementation.md | 14 + .../startup-schema-reconcile/requirements.md | 18 +- internal/database/batch_hitl_test.go | 4 +- internal/database/database.go | 1570 +---------------- internal/database/rbac.go | 141 -- internal/database/schema_legacy.go | 164 ++ internal/database/schema_reconcile.go | 288 +++ internal/database/schema_reconcile_test.go | 236 +++ internal/database/schemas/README.md | 23 + internal/database/schemas/application.sql | 862 +++++++++ internal/database/schemas/knowledge.sql | 44 + internal/database/workflow.go | 21 - internal/handler/hitl.go | 31 +- internal/handler/hitl_logs.go | 10 +- internal/handler/notification.go | 27 +- internal/knowledge/schema_migrate.go | 44 +- 18 files changed, 1669 insertions(+), 1836 deletions(-) create mode 100644 docs/feature/startup-schema-reconcile/implementation.md create mode 100644 internal/database/schema_legacy.go create mode 100644 internal/database/schema_reconcile.go create mode 100644 internal/database/schema_reconcile_test.go create mode 100644 internal/database/schemas/README.md create mode 100644 internal/database/schemas/application.sql create mode 100644 internal/database/schemas/knowledge.sql diff --git a/docs/feature/startup-schema-reconcile/design.md b/docs/feature/startup-schema-reconcile/design.md index 37ebbf66..6909c5d2 100644 --- a/docs/feature/startup-schema-reconcile/design.md +++ b/docs/feature/startup-schema-reconcile/design.md @@ -22,10 +22,14 @@ flowchart LR - 不重建表、不凭空为旧数据编造必填值。SQLite 不支持安全 ADD COLUMN 的主键、唯一约束或缺默认值的必填字段，明确报错并回滚结构补齐事务；修复定义/数据后可重新启动。 - 表级约束不自动改变。结构补齐独立于历史数据回填和既有漏洞外键语义迁移；保留外键迁移时应保留额外字段与数据，避免固定列复制造成损失。 - 全部结构补齐必须在索引及业务 SQL 之前完成；任一步失败阻止启动。现有连接保持外键配置；内存参考库用单连接并在结束时关闭。 -- 新字段仅更新 SQL 定义即可，不新增手写补字段列表。历史兼容函数可暂保留供既有测试使用，启动不调用纯字段迁移。 +- 新字段仅更新 SQL 定义即可，不新增手写补字段列表。删除纯字段迁移列表；既有测试改为验证统一结构补齐。 ## 兼容与验证 保证已有字段不被改写；默认值用于新列的已有行，未知表/字段保留。针对空库、旧库、重复启动、带引号和复杂默认值、失败回滚、非迁移列表字段、额外字段及外键数据保留做测试。已有时间回填继续保持兼容。 维护性门禁选择 adapter_extraction：从 1,822 行 database.go 提取静态 DDL 与补齐行为，旧入口薄委托；不进行无关重构。 + +## 覆盖检查与维护门禁补充 + +全仓库搜索发现处理器另有 HITL 两表与通知已读表，统一移入 application.sql，处理器仅保留数据恢复职责。HITL 文件 1,227 行（high）、通知文件 837 行（medium），只删除静态 DDL 并保留原入口/请求行为，不增加业务复杂度。最终定义覆盖 52 张表（含各数据库文件的同名表）。knowledge 的兼容入口复用同一结构引擎，不再维护字段列表。漏洞历史外键迁移仍按既有语义执行，但基于最新定义、动态复制旧列，保留额外字段、索引、触发器和依赖记录。 diff --git a/docs/feature/startup-schema-reconcile/implementation-plan.md b/docs/feature/startup-schema-reconcile/implementation-plan.md index 3e39df35..e0057ac9 100644 --- a/docs/feature/startup-schema-reconcile/implementation-plan.md +++ b/docs/feature/startup-schema-reconcile/implementation-plan.md @@ -7,7 +7,7 @@ 1. 提取现有初始化后的最终 SQL 到 internal/database/schemas，验证其覆盖所有表及索引。 2. 新增 schema_reconcile.go，通过内存参考库与事务实现通用补齐；DDL 分词用于提取准确的列定义。 3. 数据库入口改为薄委托，删除启动对纯字段迁移的依赖；保留数据回填，既有漏洞外键迁移保存额外列。 -4. 添加数据库级回归测试，运行 go test ./internal/database ./internal/knowledge 与 git diff --check。 +4. 添加数据库级回归测试，运行 go test ./internal/database ./internal/knowledge ./internal/handler ./internal/app 与 git diff --check。 5. 更新行为文档及 CHANGELOG，记录验证和提交，推送功能分支；本任务不合并或发布。 回滚：恢复代码版本；新增表和字段保留，不自动删除；升级前的历史数据不改写。验证使用临时 SQLite 文件，不操作用户实际数据库。 diff --git a/docs/feature/startup-schema-reconcile/implementation.md b/docs/feature/startup-schema-reconcile/implementation.md new file mode 100644 index 00000000..8452ca8b --- /dev/null +++ b/docs/feature/startup-schema-reconcile/implementation.md @@ -0,0 +1,14 @@ +# SQLite 启动结构补齐实现记录 + +- Phase: F4 complete +- Branch: codex/startup-schema-reconcile +- 结构来源：internal/database/schemas/*.sql，嵌入二进制。 +- 自动对比：内存 SQLite 解析最新定义，真实连接事务内创建缺表、补字段、创建缺索引。 +- 启动入口：NewDB / NewKnowledgeDB，检查失败关闭连接并返回错误。 +- 处理器中 HITL/通知表定义收敛到统一定义；请求不再临时创建通知表。 +- 删除 database/RBAC/workflow/knowledge/HITL 的纯字段迁移列表。时间、模型用量、安全拦截和审批等数据回填保留。 +- 原数据库定义文件从 1,822 行缩减到 270 行；结构与历史语义迁移各自独立。 +- 原漏洞外键迁移改为从最新定义创建目标表，动态保存额外列及其值，保留索引、触发器和依赖记录；它不是缺字段检查的依据。 +- 不删除额外表/字段，不自动修改已存在字段类型或约束。不可安全 ADD COLUMN 的定义导致补齐事务回滚，启动失败，附表/字段及原始错误。 + +实现验证：go test ./internal/database ./internal/knowledge ./internal/handler 通过；go test ./internal/app 通过；格式整理后 go test ./internal/database 再次通过。 diff --git a/docs/feature/startup-schema-reconcile/requirements.md b/docs/feature/startup-schema-reconcile/requirements.md index 8b496e98..dbf77387 100644 --- a/docs/feature/startup-schema-reconcile/requirements.md +++ b/docs/feature/startup-schema-reconcile/requirements.md @@ -16,10 +16,10 @@ | --- | --- | --- | | REQ-001 | 每次启动时对比所有应用管理表的最新定义和实际结构。 | 用户要求 | | REQ-002 | 缺表时自动创建，缺字段时自动补齐，不依赖逐字段迁移列表。 | 用户要求 | -| REQ-003 | 主数据库和知识库数据库均执行上述检查。 | 假设，待确认 | -| REQ-004 | 已有数据与额外表、字段予以保留；本次不自动删除、重命名字段或修改已有字段类型、约束。 | 假设，待确认 | -| REQ-005 | 无法安全补齐时停止启动，并报告表、字段和原因。 | 假设，待确认 | -| REQ-006 | 重复启动不重复添加表或字段，结构检查在依赖相关字段的索引创建及业务访问之前完成。 | 假设，待确认 | +| REQ-003 | 项目实际使用的 SQLite 数据库文件均执行上述检查，逐文件检查项目定义的所有表。 | 范围边界 | +| REQ-004 | 已有数据与额外表、字段予以保留；本次不自动删除、重命名字段或修改已有字段类型、约束。 | 范围边界 | +| REQ-005 | 无法安全补齐时停止启动，并报告表、字段和原因。 | 范围边界 | +| REQ-006 | 重复启动不重复添加表或字段，结构检查在依赖相关字段的索引创建及业务访问之前完成。 | 范围边界 | ## 用户场景与验收 @@ -31,23 +31,23 @@ | AC-004 | REQ-005 | 补齐操作失败时服务不开始接收请求，错误明确指出失败位置和原因。 | | AC-005 | REQ-006 | 连续启动两次成功，第二次无重复字段错误，依赖新字段的索引可正常创建。 | -## 待确认假设与范围 +## 范围与执行边界 - “所有表”指项目定义并管理的表，不包含 SQLite 内部表或其他应用自行添加的表。 - 本次需求是补齐缺表、缺字段；完整的已有类型、默认值及约束变更不纳入自动修复范围。 - SQLite 不能直接添加的字段定义应明确报错，不擅自降低约束或清空已有数据。 - 历史数据回填、外键修正等语义迁移仍需保留；它们与缺字段检查是不同职责。 -- 推荐默认：覆盖主库和知识库、保留已有数据、补齐失败阻止启动。 +- 执行范围：项目使用的 SQLite 文件、保留已有数据、补齐失败阻止启动。 ## 工作流与维护性检查 -- Workflow Gate：P1 产品契约；后端数据库初始化功能。已有启动入口、建表定义和迁移测试；缺少本功能已确认需求与设计。当前仅允许需求文档，确认后推进设计和实现。 +- Workflow Gate：P1 产品契约；后端数据库初始化功能。已有启动入口、建表定义和迁移测试；缺少本功能已确认需求与设计。需求已获明确确认，设计和实现允许推进。 - 执行范围：数据库启动初始化；不改变对外 API。 - Maintainability Gate：database.go 为 1,822 行，风险 high。采用独立模块与薄委托的 adapter_extraction，避免继续在大文件中堆积结构对比逻辑；设计与实现阶段需明确提取边界。 -- Feature Lifecycle：F0 完成仓库检查，原 main 工作区干净；专用分支 codex/startup-schema-reconcile。F1 待确认。文档为本文件；后续阶段记录设计、实现、验证及变更说明。 +- Feature Lifecycle：F0 完成仓库检查，原 main 工作区干净；专用分支 codex/startup-schema-reconcile。F1 已确认（2026-10-10 当前对话）。文档为本文件；后续阶段记录设计、实现、验证及变更说明。 ## 确认 Decision: Confirmed — 用户明确要求按启动对比、缺表补表、缺字段补字段、保留已有数据实施。 -请确认上述需求快照，或指出需要调整的需求 ID。确认后进入设计和实现。 +用户原文：“对，就按这个改，启动时自动对比 SQLite 最新结构，缺表补表、缺字段补字段，保留已有数据。” diff --git a/internal/database/batch_hitl_test.go b/internal/database/batch_hitl_test.go index 9045f317..3f536860 100644 --- a/internal/database/batch_hitl_test.go +++ b/internal/database/batch_hitl_test.go @@ -19,10 +19,10 @@ func TestBatchHITLLegacyMigration(t *testing.T) { if _, err := db.Exec("ALTER TABLE batch_task_queues DROP COLUMN hitl_policy"); err != nil { t.Fatal(err) } - if err := db.migrateBatchTaskQueuesTable(); err != nil { + if err := db.reconcileSchema(applicationSchema); err != nil { t.Fatal(err) } - if err := db.migrateBatchTaskQueuesTable(); err != nil { + if err := db.reconcileSchema(applicationSchema); err != nil { t.Fatal(err) } row, err := db.GetBatchQueue("legacy") diff --git a/internal/database/database.go b/internal/database/database.go index 8cc2f13e..97aefafc 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -191,1477 +191,23 @@ func (db *DB) SetChatUploadsDir(dir string) { // initTables 初始化数据库表 func (db *DB) initTables() error { - // 创建对话表（last_react_input / last_react_output 存「代理消息轨迹」JSON 与助手摘要，列名保留以兼容已有库） - createConversationsTable := ` - CREATE TABLE IF NOT EXISTS conversations ( - id TEXT PRIMARY KEY, - title TEXT NOT NULL, - created_at DATETIME NOT NULL, - updated_at DATETIME NOT NULL, - role_name TEXT NOT NULL DEFAULT '默认', - agent_mode TEXT NOT NULL DEFAULT 'eino_single', - last_react_input TEXT, - last_react_output TEXT - );` - - // 创建消息表 - createMessagesTable := ` - CREATE TABLE IF NOT EXISTS messages ( - id TEXT PRIMARY KEY, - conversation_id TEXT NOT NULL, - role TEXT NOT NULL, - content TEXT NOT NULL, - mcp_execution_ids TEXT, - created_at DATETIME NOT NULL, - updated_at DATETIME NOT NULL, - FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE - );` - - // 创建过程详情表 - createProcessDetailsTable := ` - CREATE TABLE IF NOT EXISTS process_details ( - id TEXT PRIMARY KEY, - message_id TEXT NOT NULL, - conversation_id TEXT NOT NULL, - event_type TEXT NOT NULL, - message TEXT, - data TEXT, - created_at DATETIME NOT NULL, - FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE, - FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE - );` - - // 创建模型 Token 用量表：process_details 负责时间线回放，本表负责结构化聚合统计。 - createModelTokenUsageTable := ` - CREATE TABLE IF NOT EXISTS model_token_usage ( - id TEXT PRIMARY KEY, - process_detail_id TEXT NOT NULL UNIQUE, - message_id TEXT NOT NULL, - conversation_id TEXT NOT NULL, - project_id TEXT, - source TEXT NOT NULL DEFAULT '', - orchestration TEXT NOT NULL DEFAULT '', - reason TEXT NOT NULL DEFAULT '', - model TEXT NOT NULL DEFAULT '', - model_calls INTEGER NOT NULL DEFAULT 0, - prompt_tokens INTEGER NOT NULL DEFAULT 0, - completion_tokens INTEGER NOT NULL DEFAULT 0, - total_tokens INTEGER NOT NULL DEFAULT 0, - cached_tokens INTEGER NOT NULL DEFAULT 0, - reasoning_tokens INTEGER NOT NULL DEFAULT 0, - created_at DATETIME NOT NULL, - updated_at DATETIME NOT NULL, - FOREIGN KEY (process_detail_id) REFERENCES process_details(id) ON DELETE CASCADE, - FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE, - FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE, - FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL - );` - - // 创建工具执行记录表 - createToolExecutionsTable := ` - CREATE TABLE IF NOT EXISTS tool_executions ( - id TEXT PRIMARY KEY, - tool_name TEXT NOT NULL, - arguments TEXT NOT NULL, - status TEXT NOT NULL, - result TEXT, - error TEXT, - start_time DATETIME NOT NULL, - end_time DATETIME, - duration_ms INTEGER, - partial_output TEXT, - partial_output_bytes INTEGER NOT NULL DEFAULT 0, - partial_output_truncated INTEGER NOT NULL DEFAULT 0, - partial_output_updated_at DATETIME, - owner_user_id TEXT, - conversation_id TEXT, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP - );` - - // 创建工具统计表 - createToolStatsTable := ` - CREATE TABLE IF NOT EXISTS tool_stats ( - tool_name TEXT PRIMARY KEY, - total_calls INTEGER NOT NULL DEFAULT 0, - success_calls INTEGER NOT NULL DEFAULT 0, - failed_calls INTEGER NOT NULL DEFAULT 0, - last_call_time DATETIME, - updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP - );` - - // 创建Skills统计表 - createSkillStatsTable := ` - CREATE TABLE IF NOT EXISTS skill_stats ( - skill_name TEXT PRIMARY KEY, - total_calls INTEGER NOT NULL DEFAULT 0, - success_calls INTEGER NOT NULL DEFAULT 0, - failed_calls INTEGER NOT NULL DEFAULT 0, - last_call_time DATETIME, - updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP - );` - - // 创建攻击链节点表 - createAttackChainNodesTable := ` - CREATE TABLE IF NOT EXISTS attack_chain_nodes ( - id TEXT PRIMARY KEY, - conversation_id TEXT NOT NULL, - node_type TEXT NOT NULL, - node_name TEXT NOT NULL, - tool_execution_id TEXT, - metadata TEXT, - risk_score INTEGER DEFAULT 0, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE, - FOREIGN KEY (tool_execution_id) REFERENCES tool_executions(id) ON DELETE SET NULL - );` - - // 创建攻击链边表 - createAttackChainEdgesTable := ` - CREATE TABLE IF NOT EXISTS attack_chain_edges ( - id TEXT PRIMARY KEY, - conversation_id TEXT NOT NULL, - source_node_id TEXT NOT NULL, - target_node_id TEXT NOT NULL, - edge_type TEXT NOT NULL, - weight INTEGER DEFAULT 1, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE, - FOREIGN KEY (source_node_id) REFERENCES attack_chain_nodes(id) ON DELETE CASCADE, - FOREIGN KEY (target_node_id) REFERENCES attack_chain_nodes(id) ON DELETE CASCADE - );` - - // 创建知识检索日志表（保留在会话数据库中，因为有外键关联） - createKnowledgeRetrievalLogsTable := ` - CREATE TABLE IF NOT EXISTS knowledge_retrieval_logs ( - id TEXT PRIMARY KEY, - conversation_id TEXT, - message_id TEXT, - query TEXT NOT NULL, - risk_type TEXT, - retrieved_items TEXT, - created_at DATETIME NOT NULL, - FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL, - FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE SET NULL - );` - - // 机器人会话绑定表（用于跨重启保持「平台+租户+用户」到 conversation 的映射） - createRobotUserSessionsTable := ` - CREATE TABLE IF NOT EXISTS robot_user_sessions ( - session_key TEXT PRIMARY KEY, - conversation_id TEXT NOT NULL, - role_name TEXT NOT NULL DEFAULT '默认', - agent_mode TEXT NOT NULL DEFAULT 'eino_single', - updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE - );` - - // 创建项目表 - createProjectsTable := ` - CREATE TABLE IF NOT EXISTS projects ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL, - description TEXT, - scope_json TEXT, - status TEXT NOT NULL DEFAULT 'active', - pinned INTEGER NOT NULL DEFAULT 0, - created_at DATETIME NOT NULL, - updated_at DATETIME NOT NULL - );` - - // 创建项目事实表（黑板） - createProjectFactsTable := ` - CREATE TABLE IF NOT EXISTS project_facts ( - id TEXT PRIMARY KEY, - project_id TEXT NOT NULL, - fact_key TEXT NOT NULL, - category TEXT NOT NULL DEFAULT 'note', - summary TEXT NOT NULL DEFAULT '', - body TEXT, - confidence TEXT NOT NULL DEFAULT 'tentative', - source_conversation_id TEXT, - source_message_id TEXT, - pinned INTEGER NOT NULL DEFAULT 0, - related_vulnerability_id TEXT, - created_at DATETIME NOT NULL, - updated_at DATETIME NOT NULL, - FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE, - UNIQUE(project_id, fact_key) - );` - - // 项目事实关系边（黑板 DAG） - createProjectFactEdgesTable := ` - CREATE TABLE IF NOT EXISTS project_fact_edges ( - id TEXT PRIMARY KEY, - project_id TEXT NOT NULL, - source_fact_key TEXT NOT NULL, - target_fact_key TEXT NOT NULL, - edge_type TEXT NOT NULL, - confidence TEXT NOT NULL DEFAULT 'tentative', - source_conversation_id TEXT, - created_at DATETIME NOT NULL, - updated_at DATETIME NOT NULL, - FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE, - UNIQUE(project_id, source_fact_key, target_fact_key, edge_type) - );` - - // 创建漏洞表 - createVulnerabilitiesTable := ` - CREATE TABLE IF NOT EXISTS vulnerabilities ( - id TEXT PRIMARY KEY, - conversation_id TEXT, - conversation_tag TEXT, - task_tag TEXT, - title TEXT NOT NULL, - description TEXT, - severity TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'open', - vulnerability_type TEXT, - target TEXT, - preconditions TEXT, - reproduction_steps TEXT, - evidence TEXT, - impact TEXT, - recommendation TEXT, - retest_notes TEXT, - retest_log TEXT NOT NULL DEFAULT '', - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - project_id TEXT, - FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL - );` - - createAssetsTable := ` - CREATE TABLE IF NOT EXISTS assets ( - id TEXT PRIMARY KEY, - dedup_key TEXT NOT NULL UNIQUE, project_id TEXT, - host TEXT NOT NULL DEFAULT '', ip TEXT NOT NULL DEFAULT '', port INTEGER NOT NULL DEFAULT 0, - domain TEXT NOT NULL DEFAULT '', protocol TEXT NOT NULL DEFAULT '', title TEXT NOT NULL DEFAULT '', - server TEXT NOT NULL DEFAULT '', country TEXT NOT NULL DEFAULT '', province TEXT NOT NULL DEFAULT '', city TEXT NOT NULL DEFAULT '', - responsible_person TEXT NOT NULL DEFAULT '', department TEXT NOT NULL DEFAULT '', business_system TEXT NOT NULL DEFAULT '', - environment TEXT NOT NULL DEFAULT '', criticality TEXT NOT NULL DEFAULT '', - source TEXT NOT NULL DEFAULT 'manual', source_query TEXT NOT NULL DEFAULT '', status TEXT NOT NULL DEFAULT 'active', - vulnerability_count INTEGER NOT NULL DEFAULT 0, risk_score INTEGER NOT NULL DEFAULT 0, risk_level TEXT NOT NULL DEFAULT 'unassessed', - tags_json TEXT NOT NULL DEFAULT '[]', first_seen_at DATETIME NOT NULL, last_seen_at DATETIME NOT NULL, - created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, owner_user_id TEXT, - FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL - );` - - createVulnerabilityAlertSubscriptionsTable := ` - CREATE TABLE IF NOT EXISTS vulnerability_alert_subscriptions ( - user_id TEXT PRIMARY KEY, - enabled INTEGER NOT NULL DEFAULT 0, - min_severity TEXT NOT NULL DEFAULT 'high', - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE - );` - createVulnerabilityAlertDeliveriesTable := ` - CREATE TABLE IF NOT EXISTS vulnerability_alert_deliveries ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - vulnerability_id TEXT NOT NULL, - user_id TEXT NOT NULL, - platform TEXT NOT NULL, - external_user_id TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'pending', - attempts INTEGER NOT NULL DEFAULT 0, - next_attempt_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - last_error TEXT NOT NULL DEFAULT '', - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - UNIQUE(vulnerability_id, platform, external_user_id), - FOREIGN KEY (vulnerability_id) REFERENCES vulnerabilities(id) ON DELETE CASCADE, - FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE - );` - - // 创建批量任务队列表 - createBatchTaskQueuesTable := ` - CREATE TABLE IF NOT EXISTS batch_task_queues ( - id TEXT PRIMARY KEY, - title TEXT, - role TEXT, - agent_mode TEXT NOT NULL DEFAULT 'eino_single', - hitl_policy TEXT NOT NULL DEFAULT '', - schedule_mode TEXT NOT NULL DEFAULT 'manual', - cron_expr TEXT, - next_run_at DATETIME, - schedule_enabled INTEGER NOT NULL DEFAULT 1, - last_schedule_trigger_at DATETIME, - last_schedule_error TEXT, - last_run_error TEXT, - project_id TEXT, - concurrency INTEGER NOT NULL DEFAULT 1, - status TEXT NOT NULL, - created_at DATETIME NOT NULL, - started_at DATETIME, - completed_at DATETIME, - current_index INTEGER NOT NULL DEFAULT 0 - );` - - // 创建批量任务表 - createBatchTasksTable := ` - CREATE TABLE IF NOT EXISTS batch_tasks ( - id TEXT PRIMARY KEY, - queue_id TEXT NOT NULL, - message TEXT NOT NULL, - conversation_id TEXT, - status TEXT NOT NULL, - started_at DATETIME, - completed_at DATETIME, - error TEXT, - result TEXT, - FOREIGN KEY (queue_id) REFERENCES batch_task_queues(id) ON DELETE CASCADE - );` - - // 创建 WebShell 连接表 - createWebshellConnectionsTable := ` - CREATE TABLE IF NOT EXISTS webshell_connections ( - id TEXT PRIMARY KEY, - project_id TEXT, - url TEXT NOT NULL, - password TEXT NOT NULL DEFAULT '', - type TEXT NOT NULL DEFAULT 'php', - method TEXT NOT NULL DEFAULT 'post', - cmd_param TEXT NOT NULL DEFAULT '', - remark TEXT NOT NULL DEFAULT '', - encoding TEXT NOT NULL DEFAULT '', - os TEXT NOT NULL DEFAULT '', - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP - );` - - // 创建 WebShell 连接扩展状态表（前端工作区/终端状态持久化） - createWebshellConnectionStatesTable := ` - CREATE TABLE IF NOT EXISTS webshell_connection_states ( - connection_id TEXT PRIMARY KEY, - state_json TEXT NOT NULL DEFAULT '{}', - updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - FOREIGN KEY (connection_id) REFERENCES webshell_connections(id) ON DELETE CASCADE - );` - - // ======================================================================== - // C2 模块（监听器 / 会话 / 任务 / 文件 / 事件 / Malleable Profile） - // ======================================================================== - createC2ListenersTable := ` - CREATE TABLE IF NOT EXISTS c2_listeners ( - id TEXT PRIMARY KEY, - project_id TEXT, - name TEXT NOT NULL, - type TEXT NOT NULL, - bind_host TEXT NOT NULL DEFAULT '127.0.0.1', - bind_port INTEGER NOT NULL, - profile_id TEXT, - encryption_key TEXT NOT NULL DEFAULT '', - implant_token TEXT NOT NULL DEFAULT '', - status TEXT NOT NULL DEFAULT 'stopped', - config_json TEXT NOT NULL DEFAULT '{}', - remark TEXT NOT NULL DEFAULT '', - owner_user_id TEXT, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - started_at DATETIME, - last_error TEXT - );` - - createC2SessionsTable := ` - CREATE TABLE IF NOT EXISTS c2_sessions ( - id TEXT PRIMARY KEY, - listener_id TEXT NOT NULL, - implant_uuid TEXT NOT NULL UNIQUE, - hostname TEXT, - username TEXT, - os TEXT, - arch TEXT, - pid INTEGER DEFAULT 0, - process_name TEXT, - is_admin INTEGER DEFAULT 0, - internal_ip TEXT, - external_ip TEXT, - user_agent TEXT, - sleep_seconds INTEGER NOT NULL DEFAULT 5, - jitter_percent INTEGER NOT NULL DEFAULT 0, - status TEXT NOT NULL DEFAULT 'active', - first_seen_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - last_check_in DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - metadata_json TEXT DEFAULT '{}', - note TEXT NOT NULL DEFAULT '', - FOREIGN KEY (listener_id) REFERENCES c2_listeners(id) ON DELETE CASCADE - );` - - createC2TasksTable := ` - CREATE TABLE IF NOT EXISTS c2_tasks ( - id TEXT PRIMARY KEY, - session_id TEXT NOT NULL, - task_type TEXT NOT NULL, - payload_json TEXT NOT NULL DEFAULT '{}', - status TEXT NOT NULL DEFAULT 'queued', - result_text TEXT, - result_blob_path TEXT, - error TEXT, - source TEXT NOT NULL DEFAULT 'manual', - conversation_id TEXT, - approval_status TEXT, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - sent_at DATETIME, - started_at DATETIME, - completed_at DATETIME, - duration_ms INTEGER DEFAULT 0, - FOREIGN KEY (session_id) REFERENCES c2_sessions(id) ON DELETE CASCADE - );` - - createC2FilesTable := ` - CREATE TABLE IF NOT EXISTS c2_files ( - id TEXT PRIMARY KEY, - session_id TEXT NOT NULL, - task_id TEXT, - direction TEXT NOT NULL, - remote_path TEXT NOT NULL, - local_path TEXT NOT NULL, - size_bytes INTEGER DEFAULT 0, - sha256 TEXT, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - FOREIGN KEY (session_id) REFERENCES c2_sessions(id) ON DELETE CASCADE - );` - - createC2EventsTable := ` - CREATE TABLE IF NOT EXISTS c2_events ( - id TEXT PRIMARY KEY, - level TEXT NOT NULL DEFAULT 'info', - category TEXT NOT NULL, - session_id TEXT, - task_id TEXT, - message TEXT NOT NULL, - data_json TEXT, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP - );` - - createAuditLogsTable := ` - CREATE TABLE IF NOT EXISTS audit_logs ( - id TEXT PRIMARY KEY, - created_at DATETIME NOT NULL, - level TEXT NOT NULL DEFAULT 'info', - category TEXT NOT NULL, - action TEXT NOT NULL, - result TEXT NOT NULL, - actor TEXT NOT NULL DEFAULT 'admin', - session_hint TEXT, - client_ip TEXT, - user_agent TEXT, - resource_type TEXT, - resource_id TEXT, - message TEXT NOT NULL, - detail_json TEXT - );` - - createC2ProfilesTable := ` - CREATE TABLE IF NOT EXISTS c2_profiles ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL UNIQUE, - user_agent TEXT, - uris_json TEXT NOT NULL DEFAULT '[]', - request_headers_json TEXT, - response_headers_json TEXT, - body_template TEXT, - jitter_min_ms INTEGER DEFAULT 0, - jitter_max_ms INTEGER DEFAULT 0, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP - );` - - createWorkflowDefinitionsTable := ` - CREATE TABLE IF NOT EXISTS workflow_definitions ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL, - description TEXT, - version INTEGER NOT NULL DEFAULT 1, - graph_json TEXT NOT NULL, - enabled INTEGER NOT NULL DEFAULT 1, - created_at DATETIME NOT NULL, - updated_at DATETIME NOT NULL - );` - - createWorkflowRunsTable := ` - CREATE TABLE IF NOT EXISTS workflow_runs ( - id TEXT PRIMARY KEY, - workflow_id TEXT NOT NULL, - workflow_version INTEGER NOT NULL DEFAULT 1, - conversation_id TEXT, - project_id TEXT, - role_id TEXT, - status TEXT NOT NULL, - input_json TEXT, - output_json TEXT, - error TEXT, - pending_hitl_node_id TEXT, - pending_hitl_json TEXT, - started_at DATETIME NOT NULL, - finished_at DATETIME, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL - );` - - createWorkflowNodeRunsTable := ` - CREATE TABLE IF NOT EXISTS workflow_node_runs ( - id TEXT PRIMARY KEY, - run_id TEXT NOT NULL, - node_id TEXT NOT NULL, - status TEXT NOT NULL, - input_json TEXT, - output_json TEXT, - error TEXT, - started_at DATETIME NOT NULL, - finished_at DATETIME, - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - FOREIGN KEY (run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE - );` - - createWorkflowPackageInspectionsTable := ` - CREATE TABLE IF NOT EXISTS workflow_package_inspections ( - id TEXT PRIMARY KEY, package_hash TEXT NOT NULL, manifest_json TEXT NOT NULL, - workflow_payload_json TEXT NOT NULL, inspection_json TEXT NOT NULL, - source_workflow_id TEXT NOT NULL, source_revision INTEGER NOT NULL, - source_content_hash TEXT NOT NULL, source_graph_hash TEXT NOT NULL, - local_conflict_state TEXT NOT NULL CHECK (local_conflict_state IN ('none','identical','id_conflict')), - local_workflow_id TEXT, local_content_hash TEXT, local_graph_hash TEXT, - created_by TEXT NOT NULL, status TEXT NOT NULL DEFAULT 'ready' CHECK (status IN ('ready','consumed','expired')), - created_at DATETIME NOT NULL, expires_at DATETIME NOT NULL, consumed_at DATETIME - );` - createWorkflowPackageImportsTable := ` - CREATE TABLE IF NOT EXISTS workflow_package_imports ( - id TEXT PRIMARY KEY, inspection_id TEXT NOT NULL, request_hash TEXT NOT NULL, - idempotency_key TEXT NOT NULL, actor_user_id TEXT NOT NULL, - action TEXT NOT NULL CHECK (action IN ('create','keep_existing','overwrite','rename')), - source_workflow_id TEXT NOT NULL, target_workflow_id TEXT NOT NULL, resulting_workflow_id TEXT, - result TEXT NOT NULL CHECK (result IN ('created','overwritten','renamed','kept_existing','skipped_identical','failed')), - error_code TEXT, error_message TEXT, created_at DATETIME NOT NULL, applied_at DATETIME, - FOREIGN KEY (inspection_id) REFERENCES workflow_package_inspections(id) - );` - - // 创建索引 - createIndexes := ` - CREATE INDEX IF NOT EXISTS idx_messages_conversation_id ON messages(conversation_id); - CREATE INDEX IF NOT EXISTS idx_conversations_updated_at ON conversations(updated_at); - CREATE INDEX IF NOT EXISTS idx_process_details_message_id ON process_details(message_id); - CREATE INDEX IF NOT EXISTS idx_process_details_conversation_id ON process_details(conversation_id); - CREATE INDEX IF NOT EXISTS idx_model_token_usage_created_at ON model_token_usage(created_at); - CREATE INDEX IF NOT EXISTS idx_model_token_usage_conversation ON model_token_usage(conversation_id); - CREATE INDEX IF NOT EXISTS idx_model_token_usage_project ON model_token_usage(project_id); - CREATE INDEX IF NOT EXISTS idx_model_token_usage_model ON model_token_usage(model); - CREATE INDEX IF NOT EXISTS idx_tool_executions_tool_name ON tool_executions(tool_name); - CREATE INDEX IF NOT EXISTS idx_tool_executions_start_time ON tool_executions(start_time); - CREATE INDEX IF NOT EXISTS idx_tool_executions_status ON tool_executions(status); - CREATE INDEX IF NOT EXISTS idx_chain_nodes_conversation ON attack_chain_nodes(conversation_id); - CREATE INDEX IF NOT EXISTS idx_chain_edges_conversation ON attack_chain_edges(conversation_id); - CREATE INDEX IF NOT EXISTS idx_chain_edges_source ON attack_chain_edges(source_node_id); - CREATE INDEX IF NOT EXISTS idx_chain_edges_target ON attack_chain_edges(target_node_id); - CREATE INDEX IF NOT EXISTS idx_knowledge_retrieval_logs_conversation ON knowledge_retrieval_logs(conversation_id); - CREATE INDEX IF NOT EXISTS idx_knowledge_retrieval_logs_message ON knowledge_retrieval_logs(message_id); - CREATE INDEX IF NOT EXISTS idx_knowledge_retrieval_logs_created_at ON knowledge_retrieval_logs(created_at); - CREATE INDEX IF NOT EXISTS idx_robot_user_sessions_updated_at ON robot_user_sessions(updated_at); - CREATE INDEX IF NOT EXISTS idx_conversations_pinned ON conversations(pinned); - CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_id ON vulnerabilities(conversation_id); - CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_tag ON vulnerabilities(conversation_tag); - CREATE INDEX IF NOT EXISTS idx_vulnerabilities_task_tag ON vulnerabilities(task_tag); - CREATE INDEX IF NOT EXISTS idx_vulnerabilities_severity ON vulnerabilities(severity); - CREATE INDEX IF NOT EXISTS idx_vulnerabilities_status ON vulnerabilities(status); - CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at); - CREATE INDEX IF NOT EXISTS idx_assets_last_seen ON assets(last_seen_at); - CREATE INDEX IF NOT EXISTS idx_assets_last_scan ON assets(last_scan_at); - CREATE INDEX IF NOT EXISTS idx_assets_ip ON assets(ip); - CREATE INDEX IF NOT EXISTS idx_assets_domain ON assets(domain); - CREATE INDEX IF NOT EXISTS idx_assets_status ON assets(status); - CREATE INDEX IF NOT EXISTS idx_assets_owner ON assets(owner_user_id); - CREATE INDEX IF NOT EXISTS idx_assets_project ON assets(project_id); - CREATE INDEX IF NOT EXISTS idx_assets_vulnerability_count ON assets(vulnerability_count); - CREATE INDEX IF NOT EXISTS idx_assets_risk_score ON assets(risk_score); - CREATE INDEX IF NOT EXISTS idx_assets_risk_level ON assets(risk_level); - CREATE INDEX IF NOT EXISTS idx_projects_status ON projects(status); - CREATE INDEX IF NOT EXISTS idx_projects_updated_at ON projects(updated_at); - CREATE INDEX IF NOT EXISTS idx_project_facts_project_id ON project_facts(project_id); - CREATE INDEX IF NOT EXISTS idx_project_facts_confidence ON project_facts(confidence); - CREATE INDEX IF NOT EXISTS idx_project_facts_related_vuln ON project_facts(related_vulnerability_id); - CREATE INDEX IF NOT EXISTS idx_project_fact_edges_project ON project_fact_edges(project_id); - CREATE INDEX IF NOT EXISTS idx_project_fact_edges_source ON project_fact_edges(project_id, source_fact_key); - CREATE INDEX IF NOT EXISTS idx_project_fact_edges_target ON project_fact_edges(project_id, target_fact_key); - CREATE INDEX IF NOT EXISTS idx_conversations_project_id ON conversations(project_id); - CREATE INDEX IF NOT EXISTS idx_vulnerabilities_project_id ON vulnerabilities(project_id); - CREATE INDEX IF NOT EXISTS idx_batch_tasks_queue_id ON batch_tasks(queue_id); - CREATE INDEX IF NOT EXISTS idx_batch_task_queues_created_at ON batch_task_queues(created_at); - CREATE INDEX IF NOT EXISTS idx_batch_task_queues_title ON batch_task_queues(title); - CREATE INDEX IF NOT EXISTS idx_webshell_connections_created_at ON webshell_connections(created_at); - CREATE INDEX IF NOT EXISTS idx_webshell_connections_project_id ON webshell_connections(project_id); - CREATE INDEX IF NOT EXISTS idx_webshell_connection_states_updated_at ON webshell_connection_states(updated_at); - CREATE INDEX IF NOT EXISTS idx_c2_listeners_created_at ON c2_listeners(created_at); - CREATE INDEX IF NOT EXISTS idx_c2_listeners_project_id ON c2_listeners(project_id); - CREATE INDEX IF NOT EXISTS idx_c2_listeners_status ON c2_listeners(status); - CREATE INDEX IF NOT EXISTS idx_c2_sessions_listener ON c2_sessions(listener_id); - CREATE INDEX IF NOT EXISTS idx_c2_sessions_status ON c2_sessions(status); - CREATE INDEX IF NOT EXISTS idx_c2_sessions_last_check_in ON c2_sessions(last_check_in); - CREATE INDEX IF NOT EXISTS idx_c2_tasks_session ON c2_tasks(session_id); - CREATE INDEX IF NOT EXISTS idx_c2_tasks_status ON c2_tasks(status); - CREATE INDEX IF NOT EXISTS idx_c2_tasks_created_at ON c2_tasks(created_at); - CREATE INDEX IF NOT EXISTS idx_c2_tasks_conversation ON c2_tasks(conversation_id); - CREATE INDEX IF NOT EXISTS idx_c2_files_session ON c2_files(session_id); - CREATE INDEX IF NOT EXISTS idx_c2_events_created_at ON c2_events(created_at); - CREATE INDEX IF NOT EXISTS idx_c2_events_category ON c2_events(category); - CREATE INDEX IF NOT EXISTS idx_c2_events_session ON c2_events(session_id); - CREATE INDEX IF NOT EXISTS idx_audit_logs_created_at ON audit_logs(created_at); - CREATE INDEX IF NOT EXISTS idx_audit_logs_category ON audit_logs(category); - CREATE INDEX IF NOT EXISTS idx_audit_logs_action ON audit_logs(action); - CREATE INDEX IF NOT EXISTS idx_audit_logs_result ON audit_logs(result); - CREATE INDEX IF NOT EXISTS idx_workflow_definitions_updated_at ON workflow_definitions(updated_at); - CREATE INDEX IF NOT EXISTS idx_workflow_definitions_enabled ON workflow_definitions(enabled); - CREATE INDEX IF NOT EXISTS idx_workflow_runs_workflow ON workflow_runs(workflow_id); - CREATE INDEX IF NOT EXISTS idx_workflow_runs_conversation ON workflow_runs(conversation_id); - CREATE INDEX IF NOT EXISTS idx_workflow_runs_status ON workflow_runs(status); - CREATE INDEX IF NOT EXISTS idx_workflow_node_runs_run ON workflow_node_runs(run_id); - CREATE INDEX IF NOT EXISTS idx_workflow_package_inspections_creator_expiry ON workflow_package_inspections(created_by, expires_at); - CREATE UNIQUE INDEX IF NOT EXISTS uq_workflow_package_imports_actor_key ON workflow_package_imports(actor_user_id, idempotency_key); - CREATE UNIQUE INDEX IF NOT EXISTS uq_workflow_package_imports_inspection_success ON workflow_package_imports(inspection_id) WHERE result IN ('created','overwritten','renamed','kept_existing','skipped_identical'); - ` - - if _, err := db.Exec(createConversationsTable); err != nil { - return fmt.Errorf("创建conversations表失败: %w", err) + if err := db.reconcileSchema(applicationSchema); err != nil { + return err } - - if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS vulnerability_retests ( - conversation_id TEXT PRIMARY KEY REFERENCES conversations(id) ON DELETE CASCADE, - source_conversation_id TEXT NOT NULL, - vulnerability_id TEXT NOT NULL, - source_project_id TEXT NOT NULL DEFAULT '', - pending_prompt TEXT NOT NULL DEFAULT '' - )`); err != nil { - return fmt.Errorf("创建复测分支表失败: %w", err) - } - - if _, err := db.Exec(createMessagesTable); err != nil { - return fmt.Errorf("创建messages表失败: %w", err) - } - - if _, err := db.Exec(createProcessDetailsTable); err != nil { - return fmt.Errorf("创建process_details表失败: %w", err) - } - - if _, err := db.Exec(createModelTokenUsageTable); err != nil { - return fmt.Errorf("创建model_token_usage表失败: %w", err) - } - - if _, err := db.Exec(createToolExecutionsTable); err != nil { - return fmt.Errorf("创建tool_executions表失败: %w", err) - } - - if _, err := db.Exec(createToolStatsTable); err != nil { - return fmt.Errorf("创建tool_stats表失败: %w", err) - } - - if _, err := db.Exec(createSkillStatsTable); err != nil { - return fmt.Errorf("创建skill_stats表失败: %w", err) - } - - if _, err := db.Exec(createAttackChainNodesTable); err != nil { - return fmt.Errorf("创建attack_chain_nodes表失败: %w", err) - } - - if _, err := db.Exec(createAttackChainEdgesTable); err != nil { - return fmt.Errorf("创建attack_chain_edges表失败: %w", err) - } - - if _, err := db.Exec(createKnowledgeRetrievalLogsTable); err != nil { - return fmt.Errorf("创建knowledge_retrieval_logs表失败: %w", err) - } - - if _, err := db.Exec(createRobotUserSessionsTable); err != nil { - return fmt.Errorf("创建robot_user_sessions表失败: %w", err) - } - if err := db.migrateRobotUserSessionsTable(); err != nil { - return fmt.Errorf("迁移robot_user_sessions表失败: %w", err) - } - - if _, err := db.Exec(createProjectsTable); err != nil { - return fmt.Errorf("创建projects表失败: %w", err) - } - - if _, err := db.Exec(createProjectFactsTable); err != nil { - return fmt.Errorf("创建project_facts表失败: %w", err) - } - - if _, err := db.Exec(createProjectFactEdgesTable); err != nil { - return fmt.Errorf("创建project_fact_edges表失败: %w", err) - } - - if _, err := db.Exec(createVulnerabilitiesTable); err != nil { - return fmt.Errorf("创建vulnerabilities表失败: %w", err) - } - if _, err := db.Exec(createAssetsTable); err != nil { - return fmt.Errorf("创建assets表失败: %w", err) - } - if err := db.migrateAssetsTable(); err != nil { - return fmt.Errorf("迁移assets表失败: %w", err) - } - - if _, err := db.Exec(createBatchTaskQueuesTable); err != nil { - return fmt.Errorf("创建batch_task_queues表失败: %w", err) - } - - if _, err := db.Exec(createBatchTasksTable); err != nil { - return fmt.Errorf("创建batch_tasks表失败: %w", err) - } - - if _, err := db.Exec(createWebshellConnectionsTable); err != nil { - return fmt.Errorf("创建webshell_connections表失败: %w", err) - } - - if _, err := db.Exec(createWebshellConnectionStatesTable); err != nil { - return fmt.Errorf("创建webshell_connection_states表失败: %w", err) - } - - if _, err := db.Exec(createAuditLogsTable); err != nil { - return fmt.Errorf("创建audit_logs表失败: %w", err) - } - - if err := db.initRBACTables(); err != nil { - return fmt.Errorf("创建RBAC表失败: %w", err) - } - if _, err := db.Exec(createVulnerabilityAlertSubscriptionsTable); err != nil { - return fmt.Errorf("创建漏洞提醒订阅表失败: %w", err) - } - if _, err := db.Exec(createVulnerabilityAlertDeliveriesTable); err != nil { - return fmt.Errorf("创建漏洞提醒投递表失败: %w", err) - } - - for tableName, ddl := range map[string]string{ - "workflow_definitions": createWorkflowDefinitionsTable, - "workflow_runs": createWorkflowRunsTable, - "workflow_node_runs": createWorkflowNodeRunsTable, - "workflow_package_inspections": createWorkflowPackageInspectionsTable, - "workflow_package_imports": createWorkflowPackageImportsTable, - } { - if _, err := db.Exec(ddl); err != nil { - return fmt.Errorf("创建%s表失败: %w", tableName, err) - } - } - - for tableName, ddl := range map[string]string{ - "c2_listeners": createC2ListenersTable, - "c2_sessions": createC2SessionsTable, - "c2_http_session_auth": `CREATE TABLE IF NOT EXISTS c2_http_session_auth ( - implant_uuid TEXT PRIMARY KEY, - listener_id TEXT NOT NULL, - token_hash TEXT NOT NULL, - FOREIGN KEY (listener_id) REFERENCES c2_listeners(id) ON DELETE CASCADE - );`, - "c2_tasks": createC2TasksTable, - "c2_files": createC2FilesTable, - "c2_events": createC2EventsTable, - "c2_profiles": createC2ProfilesTable, - } { - if _, err := db.Exec(ddl); err != nil { - return fmt.Errorf("创建%s表失败: %w", tableName, err) - } - } - - // 为已有表添加新字段（如果不存在）- 必须在创建索引之前 - if err := db.migrateConversationsTable(); err != nil { - db.logger.Warn("迁移conversations表失败", zap.Error(err)) - // 不返回错误，允许继续运行 - } - - if err := db.migrateMessagesTable(); err != nil { - db.logger.Warn("迁移messages表失败", zap.Error(err)) - // 不返回错误，允许继续运行 - } - - if err := db.migrateBatchTaskQueuesTable(); err != nil { - db.logger.Warn("迁移batch_task_queues表失败", zap.Error(err)) - // 不返回错误，允许继续运行 - } - if err := db.migrateVulnerabilitiesTable(); err != nil { - db.logger.Warn("迁移vulnerabilities表失败", zap.Error(err)) - // 不返回错误，允许继续运行 + // Data migrations are separate from schema repair. + if _, err := db.Exec("UPDATE messages SET updated_at = created_at WHERE updated_at IS NULL OR updated_at = ''"); err != nil { + return fmt.Errorf("回填消息更新时间失败: %w", err) } if err := db.migrateVulnerabilitiesConversationFK(); err != nil { - db.logger.Warn("迁移vulnerabilities会话外键失败", zap.Error(err)) + return fmt.Errorf("迁移漏洞会话外键失败: %w", err) } - - if err := db.migrateProjectsTable(); err != nil { - db.logger.Warn("迁移projects相关表失败", zap.Error(err)) - } - if err := db.dropProjectFactVersionsTable(); err != nil { - db.logger.Warn("清理project_fact_versions表失败", zap.Error(err)) - } - - if err := db.migrateWebshellConnectionsTable(); err != nil { - db.logger.Warn("迁移webshell_connections表失败", zap.Error(err)) - // 不返回错误，允许继续运行 - } - if err := db.migrateC2ListenersTable(); err != nil { - db.logger.Warn("迁移c2_listeners表失败", zap.Error(err)) - } - if err := db.migrateWorkflowRunsTable(); err != nil { - db.logger.Warn("迁移workflow_runs表失败", zap.Error(err)) - } - if err := db.migrateToolExecutionsPartialOutputColumns(); err != nil { - db.logger.Warn("迁移tool_executions partial output字段失败", zap.Error(err)) - } - if err := db.migrateRBACOwnershipColumns(); err != nil { - db.logger.Warn("迁移RBAC资源归属字段失败", zap.Error(err)) - } - - if _, err := db.Exec(createIndexes); err != nil { - return fmt.Errorf("创建索引失败: %w", err) - } - if err := db.BackfillModelTokenUsageFromProcessDetails(); err != nil { return fmt.Errorf("回填模型Token用量失败: %w", err) } - db.logger.Debug("数据库表初始化完成") + db.logger.Debug("数据库结构检查及补齐完成") return nil } -func (db *DB) migrateRobotUserSessionsTable() error { - var count int - if err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('robot_user_sessions') WHERE name='agent_mode'").Scan(&count); err != nil { - return err - } - if count == 0 { - _, err := db.Exec("ALTER TABLE robot_user_sessions ADD COLUMN agent_mode TEXT NOT NULL DEFAULT 'eino_single'") - return err - } - return nil -} - -func (db *DB) migrateToolExecutionsPartialOutputColumns() error { - for _, col := range []struct { - name string - stmt string - }{ - {"partial_output", "ALTER TABLE tool_executions ADD COLUMN partial_output TEXT"}, - {"partial_output_bytes", "ALTER TABLE tool_executions ADD COLUMN partial_output_bytes INTEGER NOT NULL DEFAULT 0"}, - {"partial_output_truncated", "ALTER TABLE tool_executions ADD COLUMN partial_output_truncated INTEGER NOT NULL DEFAULT 0"}, - {"partial_output_updated_at", "ALTER TABLE tool_executions ADD COLUMN partial_output_updated_at DATETIME"}, - } { - if err := db.addColumnIfMissing("tool_executions", col.name, col.stmt); err != nil { - return err - } - } - return nil -} - -// migrateAssetsTable keeps databases created by the first asset-management release compatible. -func (db *DB) migrateAssetsTable() error { - columns := []struct { - name string - ddl string - }{ - {"project_id", "ALTER TABLE assets ADD COLUMN project_id TEXT"}, - {"last_scan_at", "ALTER TABLE assets ADD COLUMN last_scan_at DATETIME"}, - {"last_scan_conversation_id", "ALTER TABLE assets ADD COLUMN last_scan_conversation_id TEXT NOT NULL DEFAULT ''"}, - {"last_scan_queue_id", "ALTER TABLE assets ADD COLUMN last_scan_queue_id TEXT NOT NULL DEFAULT ''"}, - {"last_scan_task_id", "ALTER TABLE assets ADD COLUMN last_scan_task_id TEXT NOT NULL DEFAULT ''"}, - {"responsible_person", "ALTER TABLE assets ADD COLUMN responsible_person TEXT NOT NULL DEFAULT ''"}, - {"department", "ALTER TABLE assets ADD COLUMN department TEXT NOT NULL DEFAULT ''"}, - {"business_system", "ALTER TABLE assets ADD COLUMN business_system TEXT NOT NULL DEFAULT ''"}, - {"environment", "ALTER TABLE assets ADD COLUMN environment TEXT NOT NULL DEFAULT ''"}, - {"criticality", "ALTER TABLE assets ADD COLUMN criticality TEXT NOT NULL DEFAULT ''"}, - {"vulnerability_count", "ALTER TABLE assets ADD COLUMN vulnerability_count INTEGER NOT NULL DEFAULT 0"}, - {"risk_score", "ALTER TABLE assets ADD COLUMN risk_score INTEGER NOT NULL DEFAULT 0"}, - {"risk_level", "ALTER TABLE assets ADD COLUMN risk_level TEXT NOT NULL DEFAULT 'unassessed'"}, - } - for _, column := range columns { - var count int - if err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('assets') WHERE name=?", column.name).Scan(&count); err != nil { - return err - } - if count == 0 { - if _, err := db.Exec(column.ddl); err != nil { - return err - } - } - } - return nil -} - -// migrateMessagesTable 迁移 messages 表，补充 updated_at 字段。 -// 语义：updated_at 表示该条消息最后一次被写入/更新的时间（例如助手占位消息在任务结束时更新正文）。 -func (db *DB) migrateMessagesTable() error { - var count int - err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('messages') WHERE name='updated_at'").Scan(&count) - if err != nil { - // 如果查询失败，尝试添加字段 - if _, addErr := db.Exec("ALTER TABLE messages ADD COLUMN updated_at DATETIME"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - return fmt.Errorf("添加 messages.updated_at 字段失败: %w", addErr) - } - } - } else if count == 0 { - if _, err := db.Exec("ALTER TABLE messages ADD COLUMN updated_at DATETIME"); err != nil { - errMsg := strings.ToLower(err.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - return fmt.Errorf("添加 messages.updated_at 字段失败: %w", err) - } - } - } - - // 回填已有数据：让 updated_at 至少等于 created_at，避免前端出现空/当前时间回退。 - _, _ = db.Exec("UPDATE messages SET updated_at = created_at WHERE updated_at IS NULL OR updated_at = ''") - - // reasoning_content：DeepSeek 思考模式 + 工具调用续跑；与 last_react_input 互补，供消息表回退路径回放 - var rcColCount int - errRC := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('messages') WHERE name='reasoning_content'").Scan(&rcColCount) - if errRC != nil { - if _, addErr := db.Exec("ALTER TABLE messages ADD COLUMN reasoning_content TEXT"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - return fmt.Errorf("添加 messages.reasoning_content 字段失败: %w", addErr) - } - } - } else if rcColCount == 0 { - if _, err := db.Exec("ALTER TABLE messages ADD COLUMN reasoning_content TEXT"); err != nil { - errMsg := strings.ToLower(err.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - return fmt.Errorf("添加 messages.reasoning_content 字段失败: %w", err) - } - } - } - return nil -} - -// migrateConversationsTable 迁移conversations表，添加新字段 -func (db *DB) migrateConversationsTable() error { - // 检查last_react_input字段是否存在 - var count int - err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversations') WHERE name='last_react_input'").Scan(&count) - if err != nil { - // 如果查询失败，尝试添加字段 - if _, addErr := db.Exec("ALTER TABLE conversations ADD COLUMN last_react_input TEXT"); addErr != nil { - // 如果字段已存在，忽略错误（SQLite错误信息可能不同） - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加last_react_input字段失败", zap.Error(addErr)) - } - } - } else if count == 0 { - // 字段不存在，添加它 - if _, err := db.Exec("ALTER TABLE conversations ADD COLUMN last_react_input TEXT"); err != nil { - db.logger.Warn("添加last_react_input字段失败", zap.Error(err)) - } - } - - // 检查last_react_output字段是否存在 - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversations') WHERE name='last_react_output'").Scan(&count) - if err != nil { - // 如果查询失败，尝试添加字段 - if _, addErr := db.Exec("ALTER TABLE conversations ADD COLUMN last_react_output TEXT"); addErr != nil { - // 如果字段已存在，忽略错误 - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加last_react_output字段失败", zap.Error(addErr)) - } - } - } else if count == 0 { - // 字段不存在，添加它 - if _, err := db.Exec("ALTER TABLE conversations ADD COLUMN last_react_output TEXT"); err != nil { - db.logger.Warn("添加last_react_output字段失败", zap.Error(err)) - } - } - - // 检查pinned字段是否存在 - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversations') WHERE name='pinned'").Scan(&count) - if err != nil { - // 如果查询失败，尝试添加字段 - if _, addErr := db.Exec("ALTER TABLE conversations ADD COLUMN pinned INTEGER DEFAULT 0"); addErr != nil { - // 如果字段已存在，忽略错误 - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加pinned字段失败", zap.Error(addErr)) - } - } - } else if count == 0 { - // 字段不存在，添加它 - if _, err := db.Exec("ALTER TABLE conversations ADD COLUMN pinned INTEGER DEFAULT 0"); err != nil { - db.logger.Warn("添加pinned字段失败", zap.Error(err)) - } - } - - // 检查 webshell_connection_id 字段是否存在（WebShell AI 助手对话关联） - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversations') WHERE name='webshell_connection_id'").Scan(&count) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE conversations ADD COLUMN webshell_connection_id TEXT"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加webshell_connection_id字段失败", zap.Error(addErr)) - } - } - } else if count == 0 { - if _, err := db.Exec("ALTER TABLE conversations ADD COLUMN webshell_connection_id TEXT"); err != nil { - db.logger.Warn("添加webshell_connection_id字段失败", zap.Error(err)) - } - } - - // 检查 role_name 字段是否存在（对话绑定的业务角色，用于历史任务切换时恢复角色上下文） - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversations') WHERE name='role_name'").Scan(&count) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE conversations ADD COLUMN role_name TEXT NOT NULL DEFAULT '默认'"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加role_name字段失败", zap.Error(addErr)) - } - } - } else if count == 0 { - if _, err := db.Exec("ALTER TABLE conversations ADD COLUMN role_name TEXT NOT NULL DEFAULT '默认'"); err != nil { - db.logger.Warn("添加role_name字段失败", zap.Error(err)) - } - } - - // 检查 agent_mode 字段是否存在（对话绑定的执行模式，用于历史任务切换时恢复对话模式） - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('conversations') WHERE name='agent_mode'").Scan(&count) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE conversations ADD COLUMN agent_mode TEXT NOT NULL DEFAULT 'eino_single'"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加agent_mode字段失败", zap.Error(addErr)) - } - } - } else if count == 0 { - if _, err := db.Exec("ALTER TABLE conversations ADD COLUMN agent_mode TEXT NOT NULL DEFAULT 'eino_single'"); err != nil { - db.logger.Warn("添加agent_mode字段失败", zap.Error(err)) - } - } - - return nil -} - -// migrateBatchTaskQueuesTable 迁移batch_task_queues表，补充新字段 -func (db *DB) migrateBatchTaskQueuesTable() error { - // 检查title字段是否存在 - var count int - err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='title'").Scan(&count) - if err != nil { - // 如果查询失败，尝试添加字段 - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN title TEXT"); addErr != nil { - // 如果字段已存在，忽略错误 - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加title字段失败", zap.Error(addErr)) - } - } - } else if count == 0 { - // 字段不存在，添加它 - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN title TEXT"); err != nil { - db.logger.Warn("添加title字段失败", zap.Error(err)) - } - } - - // 检查role字段是否存在 - var roleCount int - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='role'").Scan(&roleCount) - if err != nil { - // 如果查询失败，尝试添加字段 - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN role TEXT"); addErr != nil { - // 如果字段已存在，忽略错误 - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加role字段失败", zap.Error(addErr)) - } - } - } else if roleCount == 0 { - // 字段不存在，添加它 - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN role TEXT"); err != nil { - db.logger.Warn("添加role字段失败", zap.Error(err)) - } - } - - // 检查agent_mode字段是否存在 - var agentModeCount int - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='agent_mode'").Scan(&agentModeCount) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN agent_mode TEXT NOT NULL DEFAULT 'eino_single'"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加agent_mode字段失败", zap.Error(addErr)) - } - } - } else if agentModeCount == 0 { - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN agent_mode TEXT NOT NULL DEFAULT 'eino_single'"); err != nil { - db.logger.Warn("添加agent_mode字段失败", zap.Error(err)) - } - } - - // 检查schedule_mode字段是否存在 - var scheduleModeCount int - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='schedule_mode'").Scan(&scheduleModeCount) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN schedule_mode TEXT NOT NULL DEFAULT 'manual'"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加schedule_mode字段失败", zap.Error(addErr)) - } - } - } else if scheduleModeCount == 0 { - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN schedule_mode TEXT NOT NULL DEFAULT 'manual'"); err != nil { - db.logger.Warn("添加schedule_mode字段失败", zap.Error(err)) - } - } - - // 检查cron_expr字段是否存在 - var cronExprCount int - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='cron_expr'").Scan(&cronExprCount) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN cron_expr TEXT"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加cron_expr字段失败", zap.Error(addErr)) - } - } - } else if cronExprCount == 0 { - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN cron_expr TEXT"); err != nil { - db.logger.Warn("添加cron_expr字段失败", zap.Error(err)) - } - } - - // 检查next_run_at字段是否存在 - var nextRunAtCount int - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='next_run_at'").Scan(&nextRunAtCount) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN next_run_at DATETIME"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加next_run_at字段失败", zap.Error(addErr)) - } - } - } else if nextRunAtCount == 0 { - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN next_run_at DATETIME"); err != nil { - db.logger.Warn("添加next_run_at字段失败", zap.Error(err)) - } - } - - // schedule_enabled：0=暂停 Cron 自动调度，1=允许（手工执行不受影响） - var scheduleEnCount int - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='schedule_enabled'").Scan(&scheduleEnCount) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN schedule_enabled INTEGER NOT NULL DEFAULT 1"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加schedule_enabled字段失败", zap.Error(addErr)) - } - } - } else if scheduleEnCount == 0 { - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN schedule_enabled INTEGER NOT NULL DEFAULT 1"); err != nil { - db.logger.Warn("添加schedule_enabled字段失败", zap.Error(err)) - } - } - - var lastTrigCount int - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='last_schedule_trigger_at'").Scan(&lastTrigCount) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN last_schedule_trigger_at DATETIME"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加last_schedule_trigger_at字段失败", zap.Error(addErr)) - } - } - } else if lastTrigCount == 0 { - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN last_schedule_trigger_at DATETIME"); err != nil { - db.logger.Warn("添加last_schedule_trigger_at字段失败", zap.Error(err)) - } - } - - var lastSchedErrCount int - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='last_schedule_error'").Scan(&lastSchedErrCount) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN last_schedule_error TEXT"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加last_schedule_error字段失败", zap.Error(addErr)) - } - } - } else if lastSchedErrCount == 0 { - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN last_schedule_error TEXT"); err != nil { - db.logger.Warn("添加last_schedule_error字段失败", zap.Error(err)) - } - } - - var lastRunErrCount int - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='last_run_error'").Scan(&lastRunErrCount) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN last_run_error TEXT"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加last_run_error字段失败", zap.Error(addErr)) - } - } - } else if lastRunErrCount == 0 { - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN last_run_error TEXT"); err != nil { - db.logger.Warn("添加last_run_error字段失败", zap.Error(err)) - } - } - - var projectIDCount int - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='project_id'").Scan(&projectIDCount) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN project_id TEXT"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加batch_task_queues.project_id字段失败", zap.Error(addErr)) - } - } - } else if projectIDCount == 0 { - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN project_id TEXT"); err != nil { - db.logger.Warn("添加batch_task_queues.project_id字段失败", zap.Error(err)) - } - } - - var hitlPolicyCount int - if err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='hitl_policy'").Scan(&hitlPolicyCount); err != nil { - return fmt.Errorf("检查队列审批字段失败: %w", err) - } - if hitlPolicyCount == 0 { - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN hitl_policy TEXT NOT NULL DEFAULT ''"); err != nil { - return fmt.Errorf("添加队列审批字段失败: %w", err) - } - } - - var concurrencyCount int - err = db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('batch_task_queues') WHERE name='concurrency'").Scan(&concurrencyCount) - if err != nil { - if _, addErr := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN concurrency INTEGER NOT NULL DEFAULT 1"); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加batch_task_queues.concurrency字段失败", zap.Error(addErr)) - } - } - } else if concurrencyCount == 0 { - if _, err := db.Exec("ALTER TABLE batch_task_queues ADD COLUMN concurrency INTEGER NOT NULL DEFAULT 1"); err != nil { - db.logger.Warn("添加batch_task_queues.concurrency字段失败", zap.Error(err)) - } - } - - return nil -} - -// migrateProjectsTable 迁移 projects / conversations / vulnerabilities 的项目关联字段。 -func (db *DB) migrateProjectsTable() error { - for _, col := range []struct { - table string - name string - stmt string - }{ - {"conversations", "project_id", "ALTER TABLE conversations ADD COLUMN project_id TEXT REFERENCES projects(id) ON DELETE SET NULL"}, - {"vulnerabilities", "project_id", "ALTER TABLE vulnerabilities ADD COLUMN project_id TEXT"}, - } { - var count int - err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info(?) WHERE name=?", col.table, col.name).Scan(&count) - if err != nil { - if _, addErr := db.Exec(col.stmt); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加字段失败", zap.String("table", col.table), zap.String("field", col.name), zap.Error(addErr)) - } - } - continue - } - if count == 0 { - if _, addErr := db.Exec(col.stmt); addErr != nil { - db.logger.Warn("添加字段失败", zap.String("table", col.table), zap.String("field", col.name), zap.Error(addErr)) - } - } - } - return nil -} - -// dropProjectFactVersionsTable 移除已废弃的事实版本归档表。 -func (db *DB) dropProjectFactVersionsTable() error { - _, err := db.Exec(`DROP TABLE IF EXISTS project_fact_versions`) - return err -} - -// migrateVulnerabilitiesConversationFK 将 vulnerabilities.conversation_id 外键改为 ON DELETE SET NULL，删除对话时保留漏洞记录。 -func (db *DB) migrateVulnerabilitiesConversationFK() error { - ok, err := vulnerabilitiesConversationFKOnDeleteSetNull(db.DB) - if err != nil { - return err - } - if ok { - return nil - } - - tx, err := db.Begin() - if err != nil { - return fmt.Errorf("开启事务失败: %w", err) - } - defer func() { _ = tx.Rollback() }() - - const createNew = ` - CREATE TABLE vulnerabilities_new ( - id TEXT PRIMARY KEY, - conversation_id TEXT, - conversation_tag TEXT, - task_tag TEXT, - title TEXT NOT NULL, - description TEXT, - severity TEXT NOT NULL, - status TEXT NOT NULL DEFAULT 'open', - vulnerability_type TEXT, - target TEXT, - preconditions TEXT, - reproduction_steps TEXT, - evidence TEXT, - impact TEXT, - recommendation TEXT, - retest_notes TEXT, - retest_log TEXT NOT NULL DEFAULT '', - created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - project_id TEXT, - FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL - );` - if _, err := tx.Exec(createNew); err != nil { - return fmt.Errorf("创建 vulnerabilities_new 失败: %w", err) - } - - const copyRows = ` - INSERT INTO vulnerabilities_new ( - id, conversation_id, conversation_tag, task_tag, title, description, - severity, status, vulnerability_type, target, preconditions, reproduction_steps, - evidence, impact, recommendation, retest_notes, retest_log, - created_at, updated_at, project_id - ) - SELECT - id, conversation_id, conversation_tag, task_tag, title, description, - severity, status, vulnerability_type, target, - COALESCE(preconditions, ''), COALESCE(reproduction_steps, ''), - COALESCE(evidence, ''), impact, recommendation, COALESCE(retest_notes, ''), COALESCE(retest_log, ''), - created_at, updated_at, project_id - FROM vulnerabilities;` - if _, err := tx.Exec(copyRows); err != nil { - return fmt.Errorf("复制 vulnerabilities 数据失败: %w", err) - } - if _, err := tx.Exec(`DROP TABLE vulnerabilities`); err != nil { - return fmt.Errorf("删除旧 vulnerabilities 表失败: %w", err) - } - if _, err := tx.Exec(`ALTER TABLE vulnerabilities_new RENAME TO vulnerabilities`); err != nil { - return fmt.Errorf("重命名 vulnerabilities 表失败: %w", err) - } - - indexes := []string{ - `CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_id ON vulnerabilities(conversation_id)`, - `CREATE INDEX IF NOT EXISTS idx_vulnerabilities_conversation_tag ON vulnerabilities(conversation_tag)`, - `CREATE INDEX IF NOT EXISTS idx_vulnerabilities_task_tag ON vulnerabilities(task_tag)`, - `CREATE INDEX IF NOT EXISTS idx_vulnerabilities_severity ON vulnerabilities(severity)`, - `CREATE INDEX IF NOT EXISTS idx_vulnerabilities_status ON vulnerabilities(status)`, - `CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at)`, - `CREATE INDEX IF NOT EXISTS idx_vulnerabilities_project_id ON vulnerabilities(project_id)`, - } - for _, stmt := range indexes { - if _, err := tx.Exec(stmt); err != nil { - return fmt.Errorf("重建 vulnerabilities 索引失败: %w", err) - } - } - - if err := tx.Commit(); err != nil { - return fmt.Errorf("提交 vulnerabilities 外键迁移失败: %w", err) - } - db.logger.Info("vulnerabilities 表已迁移：删除对话时保留漏洞记录") - return nil -} - -func vulnerabilitiesConversationFKOnDeleteSetNull(db *sql.DB) (bool, error) { - rows, err := db.Query(`PRAGMA foreign_key_list(vulnerabilities)`) - if err != nil { - return false, err - } - defer rows.Close() - - found := false - for rows.Next() { - var id, seq int - var table, from, to, onUpdate, onDelete, match string - if err := rows.Scan(&id, &seq, &table, &from, &to, &onUpdate, &onDelete, &match); err != nil { - return false, err - } - if from == "conversation_id" { - found = true - if !strings.EqualFold(onDelete, "SET NULL") { - return false, nil - } - } - } - if err := rows.Err(); err != nil { - return false, err - } - return found, nil -} - -// migrateVulnerabilitiesTable 迁移 vulnerabilities 表，补充标签字段 -func (db *DB) migrateVulnerabilitiesTable() error { - columns := []struct { - name string - stmt string - }{ - {name: "conversation_tag", stmt: "ALTER TABLE vulnerabilities ADD COLUMN conversation_tag TEXT"}, - {name: "task_tag", stmt: "ALTER TABLE vulnerabilities ADD COLUMN task_tag TEXT"}, - {name: "project_id", stmt: "ALTER TABLE vulnerabilities ADD COLUMN project_id TEXT"}, - {name: "preconditions", stmt: "ALTER TABLE vulnerabilities ADD COLUMN preconditions TEXT"}, - {name: "reproduction_steps", stmt: "ALTER TABLE vulnerabilities ADD COLUMN reproduction_steps TEXT"}, - {name: "evidence", stmt: "ALTER TABLE vulnerabilities ADD COLUMN evidence TEXT"}, - {name: "retest_notes", stmt: "ALTER TABLE vulnerabilities ADD COLUMN retest_notes TEXT"}, - {name: "retest_log", stmt: "ALTER TABLE vulnerabilities ADD COLUMN retest_log TEXT NOT NULL DEFAULT ''"}, - } - - for _, col := range columns { - var count int - err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('vulnerabilities') WHERE name=?", col.name).Scan(&count) - if err != nil { - if _, addErr := db.Exec(col.stmt); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加vulnerabilities字段失败", zap.String("field", col.name), zap.Error(addErr)) - } - } - continue - } - if count == 0 { - if _, addErr := db.Exec(col.stmt); addErr != nil { - db.logger.Warn("添加vulnerabilities字段失败", zap.String("field", col.name), zap.Error(addErr)) - } - } - } - return nil -} - -// migrateWebshellConnectionsTable 迁移 webshell_connections 表，补充新字段 -func (db *DB) migrateWebshellConnectionsTable() error { - columns := []struct { - name string - stmt string - }{ - {name: "project_id", stmt: "ALTER TABLE webshell_connections ADD COLUMN project_id TEXT"}, - {name: "encoding", stmt: "ALTER TABLE webshell_connections ADD COLUMN encoding TEXT NOT NULL DEFAULT ''"}, - {name: "os", stmt: "ALTER TABLE webshell_connections ADD COLUMN os TEXT NOT NULL DEFAULT ''"}, - } - - for _, col := range columns { - var count int - err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('webshell_connections') WHERE name=?", col.name).Scan(&count) - if err != nil { - if _, addErr := db.Exec(col.stmt); addErr != nil { - errMsg := strings.ToLower(addErr.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - db.logger.Warn("添加webshell_connections字段失败", zap.String("field", col.name), zap.Error(addErr)) - } - } - continue - } - if count == 0 { - if _, addErr := db.Exec(col.stmt); addErr != nil { - db.logger.Warn("添加webshell_connections字段失败", zap.String("field", col.name), zap.Error(addErr)) - } - } - } - return nil -} - -func (db *DB) migrateC2ListenersTable() error { - return db.addColumnIfMissing("c2_listeners", "project_id", "ALTER TABLE c2_listeners ADD COLUMN project_id TEXT") -} - // NewKnowledgeDB 创建知识库数据库连接（只包含知识库相关的表） func NewKnowledgeDB(dbPath string, logger *zap.Logger) (*DB, error) { sqlDB, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_foreign_keys=1&_busy_timeout=5000&_synchronous=NORMAL") @@ -1697,108 +243,10 @@ func NewKnowledgeDB(dbPath string, logger *zap.Logger) (*DB, error) { // initKnowledgeTables 初始化知识库数据库表（只包含知识库相关的表） func (db *DB) initKnowledgeTables() error { - // 创建知识库项表 - createKnowledgeBaseItemsTable := ` - CREATE TABLE IF NOT EXISTS knowledge_base_items ( - id TEXT PRIMARY KEY, - category TEXT NOT NULL, - title TEXT NOT NULL, - file_path TEXT NOT NULL, - content TEXT, - created_at DATETIME NOT NULL, - updated_at DATETIME NOT NULL - );` - - // 创建知识库向量表 - createKnowledgeEmbeddingsTable := ` - CREATE TABLE IF NOT EXISTS knowledge_embeddings ( - id TEXT PRIMARY KEY, - item_id TEXT NOT NULL, - chunk_index INTEGER NOT NULL, - chunk_text TEXT NOT NULL, - embedding TEXT NOT NULL, - sub_indexes TEXT NOT NULL DEFAULT '', - embedding_model TEXT NOT NULL DEFAULT '', - embedding_dim INTEGER NOT NULL DEFAULT 0, - created_at DATETIME NOT NULL, - FOREIGN KEY (item_id) REFERENCES knowledge_base_items(id) ON DELETE CASCADE - );` - - // 创建知识检索日志表（在独立知识库数据库中，不使用外键约束，因为conversations和messages表可能不在这个数据库中） - createKnowledgeRetrievalLogsTable := ` - CREATE TABLE IF NOT EXISTS knowledge_retrieval_logs ( - id TEXT PRIMARY KEY, - conversation_id TEXT, - message_id TEXT, - query TEXT NOT NULL, - risk_type TEXT, - retrieved_items TEXT, - created_at DATETIME NOT NULL - );` - - // 创建索引 - createIndexes := ` - CREATE INDEX IF NOT EXISTS idx_knowledge_items_category ON knowledge_base_items(category); - CREATE INDEX IF NOT EXISTS idx_knowledge_embeddings_item_id ON knowledge_embeddings(item_id); - CREATE INDEX IF NOT EXISTS idx_knowledge_retrieval_logs_conversation ON knowledge_retrieval_logs(conversation_id); - CREATE INDEX IF NOT EXISTS idx_knowledge_retrieval_logs_message ON knowledge_retrieval_logs(message_id); - CREATE INDEX IF NOT EXISTS idx_knowledge_retrieval_logs_created_at ON knowledge_retrieval_logs(created_at); - ` - - if _, err := db.Exec(createKnowledgeBaseItemsTable); err != nil { - return fmt.Errorf("创建knowledge_base_items表失败: %w", err) - } - - if _, err := db.Exec(createKnowledgeEmbeddingsTable); err != nil { - return fmt.Errorf("创建knowledge_embeddings表失败: %w", err) - } - - if _, err := db.Exec(createKnowledgeRetrievalLogsTable); err != nil { - return fmt.Errorf("创建knowledge_retrieval_logs表失败: %w", err) - } - - if _, err := db.Exec(createIndexes); err != nil { - return fmt.Errorf("创建索引失败: %w", err) - } - - if err := db.migrateKnowledgeEmbeddingsColumns(); err != nil { - return fmt.Errorf("迁移 knowledge_embeddings 列失败: %w", err) - } - - db.logger.Info("知识库数据库表初始化完成") - return nil -} - -// migrateKnowledgeEmbeddingsColumns 为已有库补充 sub_indexes、embedding_model、embedding_dim。 -func (db *DB) migrateKnowledgeEmbeddingsColumns() error { - var n int - if err := db.QueryRow(`SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='knowledge_embeddings'`).Scan(&n); err != nil { + if err := db.reconcileSchema(knowledgeSchema); err != nil { return err } - if n == 0 { - return nil - } - migrations := []struct { - col string - stmt string - }{ - {"sub_indexes", `ALTER TABLE knowledge_embeddings ADD COLUMN sub_indexes TEXT NOT NULL DEFAULT ''`}, - {"embedding_model", `ALTER TABLE knowledge_embeddings ADD COLUMN embedding_model TEXT NOT NULL DEFAULT ''`}, - {"embedding_dim", `ALTER TABLE knowledge_embeddings ADD COLUMN embedding_dim INTEGER NOT NULL DEFAULT 0`}, - } - for _, m := range migrations { - var colCount int - q := `SELECT COUNT(*) FROM pragma_table_info('knowledge_embeddings') WHERE name = ?` - if err := db.QueryRow(q, m.col).Scan(&colCount); err != nil { - return err - } - if colCount > 0 { - continue - } - if _, err := db.Exec(m.stmt); err != nil { - return err - } - } + db.logger.Info("数据库结构检查及补齐完成") return nil } diff --git a/internal/database/rbac.go b/internal/database/rbac.go index 17db5cfa..9926ed54 100644 --- a/internal/database/rbac.go +++ b/internal/database/rbac.go @@ -87,147 +87,6 @@ type RBACAccess struct { Scope string `json:"scope"` } -func (db *DB) initRBACTables() error { - stmts := []string{ - `CREATE TABLE IF NOT EXISTS rbac_users ( - id TEXT PRIMARY KEY, - username TEXT NOT NULL UNIQUE, - display_name TEXT NOT NULL DEFAULT '', - password_hash TEXT NOT NULL, - enabled INTEGER NOT NULL DEFAULT 1, - is_builtin INTEGER NOT NULL DEFAULT 0, - created_at DATETIME NOT NULL, - updated_at DATETIME NOT NULL - );`, - `CREATE TABLE IF NOT EXISTS rbac_roles ( - id TEXT PRIMARY KEY, - name TEXT NOT NULL UNIQUE, - description TEXT NOT NULL DEFAULT '', - scope TEXT NOT NULL DEFAULT 'assigned', - is_system INTEGER NOT NULL DEFAULT 0, - created_at DATETIME NOT NULL, - updated_at DATETIME NOT NULL - );`, - `CREATE TABLE IF NOT EXISTS rbac_permissions ( - key TEXT PRIMARY KEY, - description TEXT NOT NULL DEFAULT '', - created_at DATETIME NOT NULL - );`, - `CREATE TABLE IF NOT EXISTS rbac_role_permissions ( - role_id TEXT NOT NULL, - permission_key TEXT NOT NULL, - created_at DATETIME NOT NULL, - PRIMARY KEY (role_id, permission_key), - FOREIGN KEY (role_id) REFERENCES rbac_roles(id) ON DELETE CASCADE, - FOREIGN KEY (permission_key) REFERENCES rbac_permissions(key) ON DELETE CASCADE - );`, - `CREATE TABLE IF NOT EXISTS rbac_user_roles ( - user_id TEXT NOT NULL, - role_id TEXT NOT NULL, - created_at DATETIME NOT NULL, - PRIMARY KEY (user_id, role_id), - FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE, - FOREIGN KEY (role_id) REFERENCES rbac_roles(id) ON DELETE CASCADE - );`, - `CREATE TABLE IF NOT EXISTS rbac_resource_assignments ( - id TEXT PRIMARY KEY, - user_id TEXT NOT NULL, - resource_type TEXT NOT NULL, - resource_id TEXT NOT NULL, - created_at DATETIME NOT NULL, - FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE, - UNIQUE(user_id, resource_type, resource_id) - );`, - `CREATE TABLE IF NOT EXISTS robot_user_bindings ( - id TEXT PRIMARY KEY, - platform TEXT NOT NULL, - external_user_id TEXT NOT NULL, - rbac_user_id TEXT NOT NULL, - enabled INTEGER NOT NULL DEFAULT 1, - created_at DATETIME NOT NULL, - updated_at DATETIME NOT NULL, - FOREIGN KEY (rbac_user_id) REFERENCES rbac_users(id) ON DELETE CASCADE, - UNIQUE(platform, external_user_id) - );`, - `CREATE TABLE IF NOT EXISTS robot_binding_codes ( - code_hash TEXT PRIMARY KEY, - rbac_user_id TEXT NOT NULL, - expires_at DATETIME NOT NULL, - used_at DATETIME, - created_at DATETIME NOT NULL, - FOREIGN KEY (rbac_user_id) REFERENCES rbac_users(id) ON DELETE CASCADE - );`, - `CREATE TABLE IF NOT EXISTS chat_upload_artifacts ( - relative_path TEXT PRIMARY KEY, - conversation_id TEXT NOT NULL, - owner_user_id TEXT NOT NULL, - created_at DATETIME NOT NULL, - FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE - );`, - `CREATE TABLE IF NOT EXISTS c2_payload_artifacts ( - filename TEXT PRIMARY KEY, - payload_id TEXT NOT NULL, - listener_id TEXT NOT NULL, - owner_user_id TEXT NOT NULL, - created_at DATETIME NOT NULL - );`, - `CREATE INDEX IF NOT EXISTS idx_rbac_user_roles_user ON rbac_user_roles(user_id);`, - `CREATE INDEX IF NOT EXISTS idx_rbac_role_permissions_role ON rbac_role_permissions(role_id);`, - `CREATE INDEX IF NOT EXISTS idx_rbac_assignments_user_resource ON rbac_resource_assignments(user_id, resource_type, resource_id);`, - `CREATE INDEX IF NOT EXISTS idx_rbac_assignments_resource ON rbac_resource_assignments(resource_type, resource_id);`, - `CREATE INDEX IF NOT EXISTS idx_robot_user_bindings_user ON robot_user_bindings(rbac_user_id);`, - `CREATE INDEX IF NOT EXISTS idx_robot_binding_codes_expiry ON robot_binding_codes(expires_at);`, - `CREATE INDEX IF NOT EXISTS idx_chat_upload_artifacts_conversation ON chat_upload_artifacts(conversation_id);`, - `CREATE INDEX IF NOT EXISTS idx_chat_upload_artifacts_owner ON chat_upload_artifacts(owner_user_id);`, - `CREATE INDEX IF NOT EXISTS idx_c2_payload_artifacts_listener ON c2_payload_artifacts(listener_id);`, - } - for _, stmt := range stmts { - if _, err := db.Exec(stmt); err != nil { - return err - } - } - return nil -} - -func (db *DB) migrateRBACOwnershipColumns() error { - for _, col := range []struct { - table string - name string - stmt string - }{ - {"projects", "owner_user_id", "ALTER TABLE projects ADD COLUMN owner_user_id TEXT"}, - {"conversations", "owner_user_id", "ALTER TABLE conversations ADD COLUMN owner_user_id TEXT"}, - {"vulnerabilities", "owner_user_id", "ALTER TABLE vulnerabilities ADD COLUMN owner_user_id TEXT"}, - {"webshell_connections", "owner_user_id", "ALTER TABLE webshell_connections ADD COLUMN owner_user_id TEXT"}, - {"batch_task_queues", "owner_user_id", "ALTER TABLE batch_task_queues ADD COLUMN owner_user_id TEXT"}, - {"c2_listeners", "owner_user_id", "ALTER TABLE c2_listeners ADD COLUMN owner_user_id TEXT"}, - {"tool_executions", "owner_user_id", "ALTER TABLE tool_executions ADD COLUMN owner_user_id TEXT"}, - {"tool_executions", "conversation_id", "ALTER TABLE tool_executions ADD COLUMN conversation_id TEXT"}, - } { - if err := db.addColumnIfMissing(col.table, col.name, col.stmt); err != nil { - return err - } - } - _, _ = db.Exec(`CREATE INDEX IF NOT EXISTS idx_tool_executions_owner ON tool_executions(owner_user_id)`) - _, _ = db.Exec(`CREATE INDEX IF NOT EXISTS idx_tool_executions_conversation ON tool_executions(conversation_id)`) - return nil -} - -func (db *DB) addColumnIfMissing(table, name, stmt string) error { - var count int - err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info(?) WHERE name=?", table, name).Scan(&count) - if err != nil || count == 0 { - if _, addErr := db.Exec(stmt); addErr != nil { - msg := strings.ToLower(addErr.Error()) - if !strings.Contains(msg, "duplicate column") && !strings.Contains(msg, "already exists") { - return fmt.Errorf("添加%s.%s字段失败: %w", table, name, addErr) - } - } - } - return nil -} - -// RBACNeedsAdminPassword reports whether the built-in admin account still needs an initial password. func (db *DB) RBACNeedsAdminPassword() (bool, error) { var userCount int if err := db.QueryRow(`SELECT COUNT(*) FROM rbac_users`).Scan(&userCount); err != nil { diff --git a/internal/database/schema_legacy.go b/internal/database/schema_legacy.go new file mode 100644 index 00000000..42a0ec41 --- /dev/null +++ b/internal/database/schema_legacy.go @@ -0,0 +1,164 @@ +package database + +import ( + "context" + "database/sql" + "fmt" + "strings" + + "go.uber.org/zap" +) + +// This is a historical foreign-key semantic migration, not a column repair +// list. Rebuild from the latest definition while retaining extra columns, +// indexes, triggers and dependent rows in old installations. +func (db *DB) migrateVulnerabilitiesConversationFK() error { + ok, err := vulnerabilitiesConversationFKOnDeleteSetNull(db.DB) + if err != nil || ok { + return err + } + reference, err := openSchemaReference(applicationSchema) + if err != nil { + return err + } + defer reference.Close() + var latest string + if err := reference.QueryRow(`SELECT sql FROM sqlite_master WHERE type='table' AND name='vulnerabilities'`).Scan(&latest); err != nil { + return err + } + createNew := "CREATE TABLE vulnerabilities_new " + latest[strings.Index(latest, "("):] + + // Changing foreign_keys inside a transaction is ineffective. Pin one + // connection and disable it before the transaction to avoid cascading drops. + ctx := context.Background() + conn, err := db.Conn(ctx) + if err != nil { + return err + } + defer conn.Close() + if _, err := conn.ExecContext(ctx, "PRAGMA foreign_keys=OFF"); err != nil { + return err + } + defer conn.ExecContext(ctx, "PRAGMA foreign_keys=ON") + tx, err := conn.BeginTx(ctx, nil) + if err != nil { + return err + } + defer tx.Rollback() + if _, err := tx.Exec(createNew); err != nil { + return err + } + var original string + if err := tx.QueryRow(`SELECT sql FROM sqlite_master WHERE type='table' AND name='vulnerabilities'`).Scan(&original); err != nil { + return err + } + definitions, err := columnDefinitions(original) + if err != nil { + return err + } + oldColumns, err := schemaColumns(tx, "vulnerabilities") + if err != nil { + return err + } + newColumns, err := schemaColumns(tx, "vulnerabilities_new") + if err != nil { + return err + } + present := map[string]bool{} + for _, c := range newColumns { + present[strings.ToLower(c.name)] = true + } + var names, values []string + for _, c := range oldColumns { + key := strings.ToLower(c.name) + if !present[key] { + if _, err := tx.Exec("ALTER TABLE vulnerabilities_new ADD COLUMN " + definitions[key]); err != nil { + return fmt.Errorf("保留漏洞额外字段 %s 失败: %w", c.name, err) + } + } + if c.hidden != 0 { + continue + } + quoted := quoteSchemaName(c.name) + names = append(names, quoted) + if key == "retest_log" { + values = append(values, "COALESCE("+quoted+", '')") + } else { + values = append(values, quoted) + } + } + rows, err := tx.Query(`SELECT sql FROM sqlite_master WHERE tbl_name='vulnerabilities' AND type IN ('index','trigger') AND sql IS NOT NULL ORDER BY type, name`) + if err != nil { + return err + } + var objects []string + for rows.Next() { + var ddl string + if err := rows.Scan(&ddl); err != nil { + rows.Close() + return err + } + objects = append(objects, ddl) + } + err = rows.Err() + rows.Close() + if err != nil { + return err + } + copyRows := "INSERT INTO vulnerabilities_new (" + strings.Join(names, ",") + ") SELECT " + strings.Join(values, ",") + " FROM vulnerabilities" + if _, err := tx.Exec(copyRows); err != nil { + return fmt.Errorf("保留漏洞数据失败: %w", err) + } + if _, err := tx.Exec("DROP TABLE vulnerabilities"); err != nil { + return err + } + if _, err := tx.Exec("ALTER TABLE vulnerabilities_new RENAME TO vulnerabilities"); err != nil { + return err + } + for _, ddl := range objects { + if _, err := tx.Exec(ddl); err != nil { + return err + } + } + rows, err = tx.Query("PRAGMA foreign_key_check") + if err != nil { + return err + } + violation := rows.Next() + err = rows.Err() + rows.Close() + if err != nil { + return err + } + if violation { + return fmt.Errorf("漏洞外键迁移校验失败，已保留原数据") + } + if err := tx.Commit(); err != nil { + return err + } + db.logger.Info("vulnerabilities 表已迁移：删除对话时保留漏洞记录", zap.Int("preserved_columns", len(oldColumns))) + return nil +} + +func vulnerabilitiesConversationFKOnDeleteSetNull(db *sql.DB) (bool, error) { + rows, err := db.Query(`PRAGMA foreign_key_list(vulnerabilities)`) + if err != nil { + return false, err + } + defer rows.Close() + found := false + for rows.Next() { + var id, seq int + var table, from, to, onUpdate, onDelete, match string + if err := rows.Scan(&id, &seq, &table, &from, &to, &onUpdate, &onDelete, &match); err != nil { + return false, err + } + if from == "conversation_id" { + found = true + if !strings.EqualFold(onDelete, "SET NULL") { + return false, nil + } + } + } + return found, rows.Err() +} diff --git a/internal/database/schema_reconcile.go b/internal/database/schema_reconcile.go new file mode 100644 index 00000000..a8339dd6 --- /dev/null +++ b/internal/database/schema_reconcile.go @@ -0,0 +1,288 @@ +package database + +import ( + "database/sql" + _ "embed" + "fmt" + "strings" +) + +// These definitions are the single source of truth for startup schema repair. +// +//go:embed schemas/application.sql +var applicationSchema string + +//go:embed schemas/knowledge.sql +var knowledgeSchema string + +type schemaObject struct{ name, ddl string } +type schemaColumn struct { + name string + hidden int +} +type schemaQuerier interface { + Query(string, ...any) (*sql.Rows, error) +} + +func schemaColumns(q schemaQuerier, table string) ([]schemaColumn, error) { + rows, err := q.Query(`SELECT name, hidden FROM pragma_table_xinfo(?)`, table) + if err != nil { + return nil, err + } + defer rows.Close() + var columns []schemaColumn + for rows.Next() { + var c schemaColumn + if err := rows.Scan(&c.name, &c.hidden); err != nil { + return nil, err + } + columns = append(columns, c) + } + return columns, rows.Err() +} + +func schemaObjects(q schemaQuerier, kind string) ([]schemaObject, error) { + rows, err := q.Query(`SELECT name, sql FROM sqlite_master WHERE type=? AND sql IS NOT NULL AND substr(name,1,7) != 'sqlite_' ORDER BY name`, kind) + if err != nil { + return nil, err + } + defer rows.Close() + var objects []schemaObject + for rows.Next() { + var o schemaObject + if err := rows.Scan(&o.name, &o.ddl); err != nil { + return nil, err + } + objects = append(objects, o) + } + return objects, rows.Err() +} + +func quoteSchemaName(name string) string { return `"` + strings.ReplaceAll(name, `"`, `""`) + `"` } + +// reconcileSchema only adds structure. SQLite validates the exact column +// definition; unsafe additions fail instead of weakening constraints or +// inventing values for existing records. The entire repair is transactional. +func (db *DB) reconcileSchema(schema string) error { + reference, err := openSchemaReference(schema) + if err != nil { + return err + } + defer reference.Close() + tables, err := schemaObjects(reference, "table") + if err != nil { + return err + } + indexes, err := schemaObjects(reference, "index") + if err != nil { + return err + } + tx, err := db.Begin() + if err != nil { + return err + } + defer tx.Rollback() + // Create every missing table first, so references are available before columns. + for _, table := range tables { + var count int + if err := tx.QueryRow(`SELECT count(*) FROM sqlite_master WHERE type='table' AND name=? COLLATE NOCASE`, table.name).Scan(&count); err != nil { + return err + } + if count == 0 { + if _, err := tx.Exec(table.ddl); err != nil { + return fmt.Errorf("补齐表 %s 失败: %w", table.name, err) + } + } + } + for _, table := range tables { + expected, err := schemaColumns(reference, table.name) + if err != nil { + return err + } + actual, err := schemaColumns(tx, table.name) + if err != nil { + return err + } + present := make(map[string]bool, len(actual)) + for _, c := range actual { + present[strings.ToLower(c.name)] = true + } + definitions, err := columnDefinitions(table.ddl) + if err != nil { + return fmt.Errorf("解析表 %s 字段失败: %w", table.name, err) + } + for _, c := range expected { + key := strings.ToLower(c.name) + if present[key] { + continue + } + definition, ok := definitions[key] + if !ok { + return fmt.Errorf("表 %s 字段 %s 缺少定义", table.name, c.name) + } + if _, err := tx.Exec("ALTER TABLE " + quoteSchemaName(table.name) + " ADD COLUMN " + definition); err != nil { + return fmt.Errorf("补齐字段 %s.%s 失败（保留原结构和数据，请检查字段约束及默认值）: %w", table.name, c.name, err) + } + } + } + for _, index := range indexes { + var count int + if err := tx.QueryRow(`SELECT count(*) FROM sqlite_master WHERE type='index' AND name=? COLLATE NOCASE`, index.name).Scan(&count); err != nil { + return err + } + if count == 0 { + if _, err := tx.Exec(index.ddl); err != nil { + return fmt.Errorf("补齐索引 %s 失败: %w", index.name, err) + } + } + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("提交 SQLite 结构补齐失败: %w", err) + } + return nil +} + +type schemaToken struct { + text string + start, end int + quoted bool +} + +// tokenizeSchema keeps offsets for exact SQL slices, ignoring comments and +// respecting escaped quotes. SQLite itself validates the SQL before this runs. +func tokenizeSchema(s string) ([]schemaToken, error) { + var tokens []schemaToken + for i := 0; i < len(s); { + if strings.ContainsRune(" \t\r\n", rune(s[i])) { + i++ + continue + } + if i+1 < len(s) && s[i:i+2] == "--" { + for i < len(s) && s[i] != '\n' { + i++ + } + continue + } + if i+1 < len(s) && s[i:i+2] == "/*" { + end := strings.Index(s[i+2:], "*/") + if end < 0 { + return nil, fmt.Errorf("未闭合 SQL 注释") + } + i += end + 4 + continue + } + start := i + quoted := false + if strings.ContainsRune("'\"`[", rune(s[i])) { + quoted = true + close := s[i] + if close == '[' { + close = ']' + } + i++ + closed := false + for i < len(s) { + if s[i] == close { + i++ + if close != ']' && i < len(s) && s[i] == close { + i++ + continue + } + closed = true + break + } + i++ + } + if !closed { + return nil, fmt.Errorf("未闭合 SQL 引号") + } + } else if strings.ContainsRune("(),;", rune(s[i])) { + i++ + } else { + for i < len(s) && !strings.ContainsRune(" \t\r\n(),;'\"`[", rune(s[i])) { + if i+1 < len(s) && (s[i:i+2] == "--" || s[i:i+2] == "/*") { + break + } + i++ + } + } + tokens = append(tokens, schemaToken{s[start:i], start, i, quoted}) + } + return tokens, nil +} + +func columnDefinitions(ddl string) (map[string]string, error) { + tokens, err := tokenizeSchema(ddl) + if err != nil { + return nil, err + } + definitions := make(map[string]string) + depth, start := 0, -1 + add := func(end int) { + if start < 0 || start >= end { + return + } + token := tokens[start] + if !token.quoted { + switch strings.ToUpper(token.text) { + case "CONSTRAINT", "PRIMARY", "UNIQUE", "CHECK", "FOREIGN": + return + } + } + name := token.text + if token.quoted { + close := name[len(name)-1:] + name = name[1 : len(name)-1] + if close != "]" { + name = strings.ReplaceAll(name, close+close, close) + } + } + definitions[strings.ToLower(name)] = strings.TrimSpace(ddl[token.start:tokens[end].start]) + } + for i, t := range tokens { + if t.quoted { + continue + } + switch t.text { + case "(": + depth++ + if depth == 1 { + start = i + 1 + } + case ")": + if depth == 1 { + add(i) + return definitions, nil + } + depth-- + case ",": + if depth == 1 { + add(i) + start = i + 1 + } + } + } + return nil, fmt.Errorf("缺少 CREATE TABLE 字段定义") +} + +func openSchemaReference(schema string) (*sql.DB, error) { + reference, err := sql.Open("sqlite3", ":memory:") + if err != nil { + return nil, err + } + reference.SetMaxOpenConns(1) + if _, err := reference.Exec(schema); err != nil { + reference.Close() + return nil, fmt.Errorf("解析最新 SQLite 结构失败: %w", err) + } + return reference, nil +} + +// EnsureKnowledgeSchema supports callers holding a plain SQL connection. +// The same definitions and repair engine are used by NewKnowledgeDB. +func EnsureKnowledgeSchema(conn *sql.DB) error { + if conn == nil { + return fmt.Errorf("db is nil") + } + return (&DB{DB: conn}).reconcileSchema(knowledgeSchema) +} diff --git a/internal/database/schema_reconcile_test.go b/internal/database/schema_reconcile_test.go new file mode 100644 index 00000000..3ed2cd01 --- /dev/null +++ b/internal/database/schema_reconcile_test.go @@ -0,0 +1,236 @@ +package database + +import ( + "database/sql" + "path/filepath" + "strings" + "testing" + + "go.uber.org/zap" +) + +func rawSchemaDB(t *testing.T) *DB { + t.Helper() + conn, err := sql.Open("sqlite3", ":memory:?_foreign_keys=1") + if err != nil { + t.Fatal(err) + } + conn.SetMaxOpenConns(1) + t.Cleanup(func() { conn.Close() }) + return &DB{DB: conn, logger: zap.NewNop()} +} +func schemaExec(t *testing.T, db *DB, statement string) { + t.Helper() + if _, err := db.Exec(statement); err != nil { + t.Fatal(err) + } +} + +func TestReconcileSchemaUsesDefinitionsWithoutMigrationList(t *testing.T) { + db := rawSchemaDB(t) + schemaExec(t, db, `CREATE TABLE "future's (table)" (id INTEGER PRIMARY KEY, payload TEXT, extra TEXT); + INSERT INTO "future's (table)" VALUES (7, 'old', 'keep'); + CREATE TABLE unrelated (value TEXT); INSERT INTO unrelated VALUES ('untouched');`) + latest := `CREATE TABLE "future's (table)" ( + id INTEGER PRIMARY KEY, payload TEXT, + /* punctuation , ( ) */ new_value TEXT NOT NULL DEFAULT 'a,b''c', + quota DECIMAL(10,2) DEFAULT 1.25 CHECK (quota IN (1.25, 2)), + "UNIQUE" TEXT DEFAULT 'quoted', + derived TEXT GENERATED ALWAYS AS (payload || ',x') VIRTUAL + ); + CREATE TABLE missing_table (id INTEGER PRIMARY KEY, source INTEGER REFERENCES "future's (table)"(id)); + CREATE INDEX new_index ON "future's (table)"(new_value);` + for i := 0; i < 2; i++ { + if err := db.reconcileSchema(latest); err != nil { + t.Fatal(err) + } + } + var payload, extra, value, derived, quoted string + var quota float64 + if err := db.QueryRow(`SELECT payload,extra,new_value,quota,derived,"UNIQUE" FROM "future's (table)" WHERE id=7`).Scan(&payload, &extra, &value, "a, &derived, "ed); err != nil { + t.Fatal(err) + } + if payload != "old" || extra != "keep" || value != "a,b'c" || quota != 1.25 || derived != "old,x" || quoted != "quoted" { + t.Fatalf("data/defaults: %q %q %q %v %q %q", payload, extra, value, quota, derived, quoted) + } + if err := db.QueryRow(`SELECT value FROM unrelated`).Scan(&value); err != nil || value != "untouched" { + t.Fatalf("unrelated data: %q %v", value, err) + } + schemaExec(t, db, `INSERT INTO missing_table VALUES(1,7)`) + if _, err := db.Exec(`INSERT INTO missing_table VALUES(2,999)`); err == nil { + t.Fatal("foreign key not enforced") + } +} + +func TestReconcileSchemaFailureRollsBackAllRepairs(t *testing.T) { + for _, addition := range []string{"required TEXT NOT NULL", "unique_value TEXT UNIQUE", "stamp TEXT DEFAULT CURRENT_TIMESTAMP", "stored TEXT GENERATED ALWAYS AS (payload) STORED"} { + t.Run(addition, func(t *testing.T) { + db := rawSchemaDB(t) + schemaExec(t, db, `CREATE TABLE z_existing (id INTEGER PRIMARY KEY,payload TEXT); INSERT INTO z_existing VALUES(1,'keep')`) + err := db.reconcileSchema(`CREATE TABLE a_missing (id INTEGER); CREATE TABLE z_existing (id INTEGER PRIMARY KEY,payload TEXT, safe TEXT,` + addition + `);`) + if err == nil || !strings.Contains(err.Error(), "z_existing.") { + t.Fatalf("missing contextual failure: %v", err) + } + var count int + if err := db.QueryRow(`SELECT count(*) FROM sqlite_master WHERE name='a_missing'`).Scan(&count); err != nil || count != 0 { + t.Fatalf("table repair not rolled back: %d %v", count, err) + } + if err := db.QueryRow(`SELECT count(*) FROM pragma_table_info('z_existing') WHERE name='safe'`).Scan(&count); err != nil || count != 0 { + t.Fatalf("column repair not rolled back: %d %v", count, err) + } + var value string + if err := db.QueryRow(`SELECT payload FROM z_existing WHERE id=1`).Scan(&value); err != nil || value != "keep" { + t.Fatalf("original data: %q %v", value, err) + } + }) + } +} + +func TestReconcileSchemaIndexFailureRollsBack(t *testing.T) { + db := rawSchemaDB(t) + schemaExec(t, db, `CREATE TABLE existing (value TEXT); INSERT INTO existing VALUES ('duplicate'),('duplicate')`) + err := db.reconcileSchema(`CREATE TABLE existing (value TEXT,added TEXT); CREATE UNIQUE INDEX unique_values ON existing(value)`) + if err == nil || !strings.Contains(err.Error(), "unique_values") { + t.Fatalf("index failure: %v", err) + } + var count int + if err := db.QueryRow(`SELECT count(*) FROM pragma_table_info('existing') WHERE name='added'`).Scan(&count); err != nil || count != 0 { + t.Fatalf("repair rollback: %d %v", count, err) + } +} + +func TestStartupRepairsTablesAndUnlistedColumnsPreservesData(t *testing.T) { + path := filepath.Join(t.TempDir(), "startup.db") + db, err := NewDB(path, zap.NewNop()) + if err != nil { + t.Fatal(err) + } + conv, err := db.CreateConversation("retained", ConversationCreateMeta{}) + if err != nil { + t.Fatal(err) + } + schemaExec(t, db, `ALTER TABLE conversations ADD COLUMN extension TEXT DEFAULT 'keep'; CREATE INDEX extension_index ON conversations(extension); + CREATE TABLE unrelated(value TEXT); INSERT INTO unrelated VALUES ('untouched'); + DROP TABLE c2_payload_artifacts; DROP TABLE rbac_roles; DROP TABLE notification_reads_by_user; + ALTER TABLE hitl_interrupts DROP COLUMN payload; + ALTER TABLE hitl_conversation_configs DROP COLUMN reviewer; + ALTER TABLE model_token_usage DROP COLUMN cached_tokens; + ALTER TABLE rbac_users DROP COLUMN display_name; + ALTER TABLE messages DROP COLUMN reasoning_content;`) + db.Close() + for i := 0; i < 2; i++ { + db, err = NewDB(path, zap.NewNop()) + if err != nil { + t.Fatal(err) + } + reference, err := openSchemaReference(applicationSchema) + if err != nil { + t.Fatal(err) + } + tables, err := schemaObjects(reference, "table") + if err != nil { + t.Fatal(err) + } + for _, table := range tables { + expected, err := schemaColumns(reference, table.name) + if err != nil { + t.Fatal(err) + } + actual, err := schemaColumns(db, table.name) + if err != nil { + t.Fatal(err) + } + present := map[string]bool{} + for _, c := range actual { + present[c.name] = true + } + for _, c := range expected { + if !present[c.name] { + t.Errorf("missing %s.%s", table.name, c.name) + } + } + } + reference.Close() + var title, extra string + if err := db.QueryRow(`SELECT title,extension FROM conversations WHERE id=?`, conv.ID).Scan(&title, &extra); err != nil || title != "retained" || extra != "keep" { + t.Fatalf("data: %q %q %v", title, extra, err) + } + var value string + if err := db.QueryRow(`SELECT value FROM unrelated`).Scan(&value); err != nil || value != "untouched" { + t.Fatalf("extra table: %q %v", value, err) + } + db.Close() + } +} + +func TestOtherSQLiteFileStartupRepairsSchema(t *testing.T) { + path := filepath.Join(t.TempDir(), "other.db") + db, err := NewKnowledgeDB(path, zap.NewNop()) + if err != nil { + t.Fatal(err) + } + schemaExec(t, db, `INSERT INTO knowledge_base_items (id,category,title,file_path,created_at,updated_at) VALUES('keep','cat','title','path','2026-10-10','2026-10-10'); + ALTER TABLE knowledge_base_items DROP COLUMN content; + ALTER TABLE knowledge_embeddings DROP COLUMN embedding_model; + DROP TABLE knowledge_retrieval_logs;`) + db.Close() + for i := 0; i < 2; i++ { + db, err = NewKnowledgeDB(path, zap.NewNop()) + if err != nil { + t.Fatal(err) + } + var title string + if err := db.QueryRow(`SELECT title FROM knowledge_base_items WHERE id='keep'`).Scan(&title); err != nil || title != "title" { + t.Fatalf("data: %q %v", title, err) + } + schemaExec(t, db, `SELECT content FROM knowledge_base_items; SELECT embedding_model FROM knowledge_embeddings; SELECT * FROM knowledge_retrieval_logs`) + db.Close() + } +} + +func TestHistoricalFKMigrationPreservesExtraAndDependentData(t *testing.T) { + db := rawSchemaDB(t) + legacy := strings.Replace(applicationSchema, `FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL`, `FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE`, 1) + schemaExec(t, db, legacy) + conv, err := db.CreateConversation("source", ConversationCreateMeta{}) + if err != nil { + t.Fatal(err) + } + vuln, err := db.CreateVulnerability(&Vulnerability{ConversationID: conv.ID, Title: "keep", Severity: "high", Status: "open"}) + if err != nil { + t.Fatal(err) + } + schemaExec(t, db, `ALTER TABLE vulnerabilities ADD COLUMN extension TEXT DEFAULT 'keep'; CREATE INDEX vuln_extra_index ON vulnerabilities(extension); + CREATE TABLE dependent(id TEXT REFERENCES vulnerabilities(id) ON DELETE CASCADE); + CREATE TABLE trigger_log(value TEXT); + CREATE TRIGGER vuln_custom AFTER UPDATE ON vulnerabilities BEGIN INSERT INTO trigger_log VALUES (NEW.extension); END;`) + if _, err := db.Exec(`INSERT INTO dependent VALUES(?)`, vuln.ID); err != nil { + t.Fatal(err) + } + if err := db.initTables(); err != nil { + t.Fatal(err) + } + var value string + if err := db.QueryRow(`SELECT extension FROM vulnerabilities WHERE id=?`, vuln.ID).Scan(&value); err != nil || value != "keep" { + t.Fatalf("extra data: %q %v", value, err) + } + var count int + if err := db.QueryRow(`SELECT count(*) FROM dependent`).Scan(&count); err != nil || count != 1 { + t.Fatalf("dependent rows: %d %v", count, err) + } + if _, err := db.Exec(`UPDATE vulnerabilities SET extension='triggered' WHERE id=?`, vuln.ID); err != nil { + t.Fatal(err) + } + if err := db.QueryRow(`SELECT value FROM trigger_log`).Scan(&value); err != nil || value != "triggered" { + t.Fatalf("trigger: %q %v", value, err) + } + if err := db.QueryRow(`PRAGMA foreign_keys`).Scan(&count); err != nil || count != 1 { + t.Fatalf("foreign keys: %d %v", count, err) + } + if err := db.DeleteConversation(conv.ID); err != nil { + t.Fatal(err) + } + if _, err := db.GetVulnerability(vuln.ID); err != nil { + t.Fatal(err) + } +} diff --git a/internal/database/schemas/README.md b/internal/database/schemas/README.md new file mode 100644 index 00000000..4a2338fb --- /dev/null +++ b/internal/database/schemas/README.md @@ -0,0 +1,23 @@ +# SQLite latest schema definitions + +These SQL files are embedded into the server binary and are the authoritative +source for the tables and columns managed by the application. When adding a +table or column, update its definition here; do not add a separate list of +`ALTER TABLE ADD COLUMN` migrations in startup or handlers. + +Startup opens an isolated in-memory reference database with the latest SQL, +compares it with the actual SQLite file, and adds missing tables and columns in +one transaction before creating missing indexes and running business queries. +Existing rows, extra tables and extra columns remain in place. Repeated startup +is safe. The process does not change types or constraints of existing columns. + +Column additions use their exact definition, including defaults. SQLite rejects +some additions (for example PRIMARY KEY, UNIQUE, STORED generated columns, or a +required column without a suitable default on a populated table). Such failures +roll back structure repair and stop startup with the table/column and cause. +Resolve the definition or explicitly plan a data migration, then restart; the +repair engine does not invent required values or weaken constraints. + +Historical data backfills and foreign-key semantic changes are separate from +missing-table/column repair. Their behavior is documented in +`docs/feature/startup-schema-reconcile/`. diff --git a/internal/database/schemas/application.sql b/internal/database/schemas/application.sql new file mode 100644 index 00000000..2bba4c23 --- /dev/null +++ b/internal/database/schemas/application.sql @@ -0,0 +1,862 @@ +-- Latest SQLite structure. Update these definitions when adding tables or columns. + +CREATE TABLE assets ( + id TEXT PRIMARY KEY, + dedup_key TEXT NOT NULL UNIQUE, + project_id TEXT, + host TEXT NOT NULL DEFAULT '', + ip TEXT NOT NULL DEFAULT '', + port INTEGER NOT NULL DEFAULT 0, + domain TEXT NOT NULL DEFAULT '', + protocol TEXT NOT NULL DEFAULT '', + title TEXT NOT NULL DEFAULT '', + server TEXT NOT NULL DEFAULT '', + country TEXT NOT NULL DEFAULT '', + province TEXT NOT NULL DEFAULT '', + city TEXT NOT NULL DEFAULT '', + responsible_person TEXT NOT NULL DEFAULT '', + department TEXT NOT NULL DEFAULT '', + business_system TEXT NOT NULL DEFAULT '', + environment TEXT NOT NULL DEFAULT '', + criticality TEXT NOT NULL DEFAULT '', + source TEXT NOT NULL DEFAULT 'manual', + source_query TEXT NOT NULL DEFAULT '', + status TEXT NOT NULL DEFAULT 'active', + vulnerability_count INTEGER NOT NULL DEFAULT 0, + risk_score INTEGER NOT NULL DEFAULT 0, + risk_level TEXT NOT NULL DEFAULT 'unassessed', + tags_json TEXT NOT NULL DEFAULT '[]', + first_seen_at DATETIME NOT NULL, + last_seen_at DATETIME NOT NULL, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL, + owner_user_id TEXT, + last_scan_at DATETIME, + last_scan_conversation_id TEXT NOT NULL DEFAULT '', + last_scan_queue_id TEXT NOT NULL DEFAULT '', + last_scan_task_id TEXT NOT NULL DEFAULT '', + FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL +); + +CREATE TABLE attack_chain_edges ( + id TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL, + source_node_id TEXT NOT NULL, + target_node_id TEXT NOT NULL, + edge_type TEXT NOT NULL, + weight INTEGER DEFAULT 1, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE, + FOREIGN KEY (source_node_id) REFERENCES attack_chain_nodes(id) ON DELETE CASCADE, + FOREIGN KEY (target_node_id) REFERENCES attack_chain_nodes(id) ON DELETE CASCADE +); + +CREATE TABLE attack_chain_nodes ( + id TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL, + node_type TEXT NOT NULL, + node_name TEXT NOT NULL, + tool_execution_id TEXT, + metadata TEXT, + risk_score INTEGER DEFAULT 0, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE, + FOREIGN KEY (tool_execution_id) REFERENCES tool_executions(id) ON DELETE SET NULL +); + +CREATE TABLE audit_logs ( + id TEXT PRIMARY KEY, + created_at DATETIME NOT NULL, + level TEXT NOT NULL DEFAULT 'info', + category TEXT NOT NULL, + action TEXT NOT NULL, + result TEXT NOT NULL, + actor TEXT NOT NULL DEFAULT 'admin', + session_hint TEXT, + client_ip TEXT, + user_agent TEXT, + resource_type TEXT, + resource_id TEXT, + message TEXT NOT NULL, + detail_json TEXT +); + +CREATE TABLE batch_task_queues ( + id TEXT PRIMARY KEY, + title TEXT, + role TEXT, + agent_mode TEXT NOT NULL DEFAULT 'eino_single', + hitl_policy TEXT NOT NULL DEFAULT '', + schedule_mode TEXT NOT NULL DEFAULT 'manual', + cron_expr TEXT, + next_run_at DATETIME, + schedule_enabled INTEGER NOT NULL DEFAULT 1, + last_schedule_trigger_at DATETIME, + last_schedule_error TEXT, + last_run_error TEXT, + project_id TEXT, + concurrency INTEGER NOT NULL DEFAULT 1, + status TEXT NOT NULL, + created_at DATETIME NOT NULL, + started_at DATETIME, + completed_at DATETIME, + current_index INTEGER NOT NULL DEFAULT 0, + owner_user_id TEXT +); + +CREATE TABLE batch_tasks ( + id TEXT PRIMARY KEY, + queue_id TEXT NOT NULL, + message TEXT NOT NULL, + conversation_id TEXT, + status TEXT NOT NULL, + started_at DATETIME, + completed_at DATETIME, + error TEXT, + result TEXT, + FOREIGN KEY (queue_id) REFERENCES batch_task_queues(id) ON DELETE CASCADE +); + +CREATE TABLE c2_events ( + id TEXT PRIMARY KEY, + level TEXT NOT NULL DEFAULT 'info', + category TEXT NOT NULL, + session_id TEXT, + task_id TEXT, + message TEXT NOT NULL, + data_json TEXT, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE TABLE c2_files ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + task_id TEXT, + direction TEXT NOT NULL, + remote_path TEXT NOT NULL, + local_path TEXT NOT NULL, + size_bytes INTEGER DEFAULT 0, + sha256 TEXT, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (session_id) REFERENCES c2_sessions(id) ON DELETE CASCADE +); + +CREATE TABLE c2_http_session_auth ( + implant_uuid TEXT PRIMARY KEY, + listener_id TEXT NOT NULL, + token_hash TEXT NOT NULL, + FOREIGN KEY (listener_id) REFERENCES c2_listeners(id) ON DELETE CASCADE +); + +CREATE TABLE c2_listeners ( + id TEXT PRIMARY KEY, + project_id TEXT, + name TEXT NOT NULL, + type TEXT NOT NULL, + bind_host TEXT NOT NULL DEFAULT '127.0.0.1', + bind_port INTEGER NOT NULL, + profile_id TEXT, + encryption_key TEXT NOT NULL DEFAULT '', + implant_token TEXT NOT NULL DEFAULT '', + status TEXT NOT NULL DEFAULT 'stopped', + config_json TEXT NOT NULL DEFAULT '{}', + remark TEXT NOT NULL DEFAULT '', + owner_user_id TEXT, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + started_at DATETIME, + last_error TEXT +); + +CREATE TABLE c2_payload_artifacts ( + filename TEXT PRIMARY KEY, + payload_id TEXT NOT NULL, + listener_id TEXT NOT NULL, + owner_user_id TEXT NOT NULL, + created_at DATETIME NOT NULL +); + +CREATE TABLE c2_profiles ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL UNIQUE, + user_agent TEXT, + uris_json TEXT NOT NULL DEFAULT '[]', + request_headers_json TEXT, + response_headers_json TEXT, + body_template TEXT, + jitter_min_ms INTEGER DEFAULT 0, + jitter_max_ms INTEGER DEFAULT 0, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE TABLE c2_sessions ( + id TEXT PRIMARY KEY, + listener_id TEXT NOT NULL, + implant_uuid TEXT NOT NULL UNIQUE, + hostname TEXT, + username TEXT, + os TEXT, + arch TEXT, + pid INTEGER DEFAULT 0, + process_name TEXT, + is_admin INTEGER DEFAULT 0, + internal_ip TEXT, + external_ip TEXT, + user_agent TEXT, + sleep_seconds INTEGER NOT NULL DEFAULT 5, + jitter_percent INTEGER NOT NULL DEFAULT 0, + status TEXT NOT NULL DEFAULT 'active', + first_seen_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + last_check_in DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + metadata_json TEXT DEFAULT '{}', + note TEXT NOT NULL DEFAULT '', + FOREIGN KEY (listener_id) REFERENCES c2_listeners(id) ON DELETE CASCADE +); + +CREATE TABLE c2_tasks ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + task_type TEXT NOT NULL, + payload_json TEXT NOT NULL DEFAULT '{}', + status TEXT NOT NULL DEFAULT 'queued', + result_text TEXT, + result_blob_path TEXT, + error TEXT, + source TEXT NOT NULL DEFAULT 'manual', + conversation_id TEXT, + approval_status TEXT, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + sent_at DATETIME, + started_at DATETIME, + completed_at DATETIME, + duration_ms INTEGER DEFAULT 0, + FOREIGN KEY (session_id) REFERENCES c2_sessions(id) ON DELETE CASCADE +); + +CREATE TABLE chat_upload_artifacts ( + relative_path TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL, + owner_user_id TEXT NOT NULL, + created_at DATETIME NOT NULL, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE +); + +CREATE TABLE conversations ( + id TEXT PRIMARY KEY, + title TEXT NOT NULL, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL, + role_name TEXT NOT NULL DEFAULT '默认', + agent_mode TEXT NOT NULL DEFAULT 'eino_single', + last_react_input TEXT, + last_react_output TEXT, + pinned INTEGER DEFAULT 0, + webshell_connection_id TEXT, + project_id TEXT REFERENCES projects(id) ON DELETE SET NULL, + owner_user_id TEXT +); + +CREATE TABLE hitl_conversation_configs ( + conversation_id TEXT PRIMARY KEY, + enabled INTEGER NOT NULL DEFAULT 0, + mode TEXT NOT NULL DEFAULT 'off', + sensitive_tools TEXT NOT NULL DEFAULT '[]', + timeout_seconds INTEGER NOT NULL DEFAULT 0, + updated_at DATETIME NOT NULL, + reviewer TEXT NOT NULL DEFAULT 'human' +); + +CREATE TABLE hitl_interrupts ( + id TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL, + message_id TEXT, + mode TEXT NOT NULL, + tool_name TEXT NOT NULL, + tool_call_id TEXT, + payload TEXT, + status TEXT NOT NULL, + reviewer TEXT NOT NULL DEFAULT 'human', + decision TEXT, + decision_comment TEXT, + created_at DATETIME NOT NULL, + decided_at DATETIME, + decided_by TEXT NOT NULL DEFAULT 'human' +); + +CREATE TABLE knowledge_retrieval_logs ( + id TEXT PRIMARY KEY, + conversation_id TEXT, + message_id TEXT, + query TEXT NOT NULL, + risk_type TEXT, + retrieved_items TEXT, + created_at DATETIME NOT NULL, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL, + FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE SET NULL +); + +CREATE TABLE messages ( + id TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL, + role TEXT NOT NULL, + content TEXT NOT NULL, + mcp_execution_ids TEXT, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL, + reasoning_content TEXT, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE +); + +CREATE TABLE model_token_usage ( + id TEXT PRIMARY KEY, + process_detail_id TEXT NOT NULL UNIQUE, + message_id TEXT NOT NULL, + conversation_id TEXT NOT NULL, + project_id TEXT, + source TEXT NOT NULL DEFAULT '', + orchestration TEXT NOT NULL DEFAULT '', + reason TEXT NOT NULL DEFAULT '', + model TEXT NOT NULL DEFAULT '', + model_calls INTEGER NOT NULL DEFAULT 0, + prompt_tokens INTEGER NOT NULL DEFAULT 0, + completion_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + cached_tokens INTEGER NOT NULL DEFAULT 0, + reasoning_tokens INTEGER NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL, + FOREIGN KEY (process_detail_id) REFERENCES process_details(id) ON DELETE CASCADE, + FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE, + FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE SET NULL +); + +CREATE TABLE notification_reads_by_user ( + user_id TEXT NOT NULL, + event_id TEXT NOT NULL, + read_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY(user_id, event_id) +); + +CREATE TABLE process_details ( + id TEXT PRIMARY KEY, + message_id TEXT NOT NULL, + conversation_id TEXT NOT NULL, + event_type TEXT NOT NULL, + message TEXT, + data TEXT, + created_at DATETIME NOT NULL, + FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE +); + +CREATE TABLE project_fact_edges ( + id TEXT PRIMARY KEY, + project_id TEXT NOT NULL, + source_fact_key TEXT NOT NULL, + target_fact_key TEXT NOT NULL, + edge_type TEXT NOT NULL, + confidence TEXT NOT NULL DEFAULT 'tentative', + source_conversation_id TEXT, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL, + FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE, + UNIQUE(project_id, source_fact_key, target_fact_key, edge_type) +); + +CREATE TABLE project_facts ( + id TEXT PRIMARY KEY, + project_id TEXT NOT NULL, + fact_key TEXT NOT NULL, + category TEXT NOT NULL DEFAULT 'note', + summary TEXT NOT NULL DEFAULT '', + body TEXT, + confidence TEXT NOT NULL DEFAULT 'tentative', + source_conversation_id TEXT, + source_message_id TEXT, + pinned INTEGER NOT NULL DEFAULT 0, + related_vulnerability_id TEXT, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL, + FOREIGN KEY (project_id) REFERENCES projects(id) ON DELETE CASCADE, + UNIQUE(project_id, fact_key) +); + +CREATE TABLE projects ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + description TEXT, + scope_json TEXT, + status TEXT NOT NULL DEFAULT 'active', + pinned INTEGER NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL, + owner_user_id TEXT +); + +CREATE TABLE rbac_permissions ( + key TEXT PRIMARY KEY, + description TEXT NOT NULL DEFAULT '', + created_at DATETIME NOT NULL +); + +CREATE TABLE rbac_resource_assignments ( + id TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + resource_type TEXT NOT NULL, + resource_id TEXT NOT NULL, + created_at DATETIME NOT NULL, + FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE, + UNIQUE(user_id, resource_type, resource_id) +); + +CREATE TABLE rbac_role_permissions ( + role_id TEXT NOT NULL, + permission_key TEXT NOT NULL, + created_at DATETIME NOT NULL, + PRIMARY KEY (role_id, permission_key), + FOREIGN KEY (role_id) REFERENCES rbac_roles(id) ON DELETE CASCADE, + FOREIGN KEY (permission_key) REFERENCES rbac_permissions(key) ON DELETE CASCADE +); + +CREATE TABLE rbac_roles ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL UNIQUE, + description TEXT NOT NULL DEFAULT '', + scope TEXT NOT NULL DEFAULT 'assigned', + is_system INTEGER NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL +); + +CREATE TABLE rbac_user_roles ( + user_id TEXT NOT NULL, + role_id TEXT NOT NULL, + created_at DATETIME NOT NULL, + PRIMARY KEY (user_id, role_id), + FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE, + FOREIGN KEY (role_id) REFERENCES rbac_roles(id) ON DELETE CASCADE +); + +CREATE TABLE rbac_users ( + id TEXT PRIMARY KEY, + username TEXT NOT NULL UNIQUE, + display_name TEXT NOT NULL DEFAULT '', + password_hash TEXT NOT NULL, + enabled INTEGER NOT NULL DEFAULT 1, + is_builtin INTEGER NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL +); + +CREATE TABLE robot_binding_codes ( + code_hash TEXT PRIMARY KEY, + rbac_user_id TEXT NOT NULL, + expires_at DATETIME NOT NULL, + used_at DATETIME, + created_at DATETIME NOT NULL, + FOREIGN KEY (rbac_user_id) REFERENCES rbac_users(id) ON DELETE CASCADE +); + +CREATE TABLE robot_user_bindings ( + id TEXT PRIMARY KEY, + platform TEXT NOT NULL, + external_user_id TEXT NOT NULL, + rbac_user_id TEXT NOT NULL, + enabled INTEGER NOT NULL DEFAULT 1, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL, + FOREIGN KEY (rbac_user_id) REFERENCES rbac_users(id) ON DELETE CASCADE, + UNIQUE(platform, external_user_id) +); + +CREATE TABLE robot_user_sessions ( + session_key TEXT PRIMARY KEY, + conversation_id TEXT NOT NULL, + role_name TEXT NOT NULL DEFAULT '默认', + agent_mode TEXT NOT NULL DEFAULT 'eino_single', + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE +); + +CREATE TABLE skill_stats ( + skill_name TEXT PRIMARY KEY, + total_calls INTEGER NOT NULL DEFAULT 0, + success_calls INTEGER NOT NULL DEFAULT 0, + failed_calls INTEGER NOT NULL DEFAULT 0, + last_call_time DATETIME, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE TABLE tool_executions ( + id TEXT PRIMARY KEY, + tool_name TEXT NOT NULL, + arguments TEXT NOT NULL, + status TEXT NOT NULL, + result TEXT, + error TEXT, + start_time DATETIME NOT NULL, + end_time DATETIME, + duration_ms INTEGER, + partial_output TEXT, + partial_output_bytes INTEGER NOT NULL DEFAULT 0, + partial_output_truncated INTEGER NOT NULL DEFAULT 0, + partial_output_updated_at DATETIME, + owner_user_id TEXT, + conversation_id TEXT, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE TABLE tool_stats ( + tool_name TEXT PRIMARY KEY, + total_calls INTEGER NOT NULL DEFAULT 0, + success_calls INTEGER NOT NULL DEFAULT 0, + failed_calls INTEGER NOT NULL DEFAULT 0, + last_call_time DATETIME, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE TABLE vulnerabilities ( + id TEXT PRIMARY KEY, + conversation_id TEXT, + conversation_tag TEXT, + task_tag TEXT, + title TEXT NOT NULL, + description TEXT, + severity TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'open', + vulnerability_type TEXT, + target TEXT, + preconditions TEXT, + reproduction_steps TEXT, + evidence TEXT, + impact TEXT, + recommendation TEXT, + retest_notes TEXT, + retest_log TEXT NOT NULL DEFAULT '', + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + project_id TEXT, + owner_user_id TEXT, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL +); + +CREATE TABLE vulnerability_alert_deliveries ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + vulnerability_id TEXT NOT NULL, + user_id TEXT NOT NULL, + platform TEXT NOT NULL, + external_user_id TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + attempts INTEGER NOT NULL DEFAULT 0, + next_attempt_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + last_error TEXT NOT NULL DEFAULT '', + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + UNIQUE(vulnerability_id, platform, external_user_id), + FOREIGN KEY (vulnerability_id) REFERENCES vulnerabilities(id) ON DELETE CASCADE, + FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE +); + +CREATE TABLE vulnerability_alert_subscriptions ( + user_id TEXT PRIMARY KEY, + enabled INTEGER NOT NULL DEFAULT 0, + min_severity TEXT NOT NULL DEFAULT 'high', + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (user_id) REFERENCES rbac_users(id) ON DELETE CASCADE +); + +CREATE TABLE vulnerability_retests ( + conversation_id TEXT PRIMARY KEY REFERENCES conversations(id) ON DELETE CASCADE, + source_conversation_id TEXT NOT NULL, + vulnerability_id TEXT NOT NULL, + source_project_id TEXT NOT NULL DEFAULT '', + pending_prompt TEXT NOT NULL DEFAULT '' +); + +CREATE TABLE webshell_connection_states ( + connection_id TEXT PRIMARY KEY, + state_json TEXT NOT NULL DEFAULT '{}', + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (connection_id) REFERENCES webshell_connections(id) ON DELETE CASCADE +); + +CREATE TABLE webshell_connections ( + id TEXT PRIMARY KEY, + project_id TEXT, + url TEXT NOT NULL, + password TEXT NOT NULL DEFAULT '', + type TEXT NOT NULL DEFAULT 'php', + method TEXT NOT NULL DEFAULT 'post', + cmd_param TEXT NOT NULL DEFAULT '', + remark TEXT NOT NULL DEFAULT '', + encoding TEXT NOT NULL DEFAULT '', + os TEXT NOT NULL DEFAULT '', + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + owner_user_id TEXT +); + +CREATE TABLE workflow_definitions ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + description TEXT, + version INTEGER NOT NULL DEFAULT 1, + graph_json TEXT NOT NULL, + enabled INTEGER NOT NULL DEFAULT 1, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL +); + +CREATE TABLE workflow_node_runs ( + id TEXT PRIMARY KEY, + run_id TEXT NOT NULL, + node_id TEXT NOT NULL, + status TEXT NOT NULL, + input_json TEXT, + output_json TEXT, + error TEXT, + started_at DATETIME NOT NULL, + finished_at DATETIME, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE +); + +CREATE TABLE workflow_package_imports ( + id TEXT PRIMARY KEY, + inspection_id TEXT NOT NULL, + request_hash TEXT NOT NULL, + idempotency_key TEXT NOT NULL, + actor_user_id TEXT NOT NULL, + action TEXT NOT NULL CHECK (action IN ('create','keep_existing','overwrite','rename')), + source_workflow_id TEXT NOT NULL, + target_workflow_id TEXT NOT NULL, + resulting_workflow_id TEXT, + result TEXT NOT NULL CHECK (result IN ('created','overwritten','renamed','kept_existing','skipped_identical','failed')), + error_code TEXT, + error_message TEXT, + created_at DATETIME NOT NULL, + applied_at DATETIME, + FOREIGN KEY (inspection_id) REFERENCES workflow_package_inspections(id) +); + +CREATE TABLE workflow_package_inspections ( + id TEXT PRIMARY KEY, + package_hash TEXT NOT NULL, + manifest_json TEXT NOT NULL, + workflow_payload_json TEXT NOT NULL, + inspection_json TEXT NOT NULL, + source_workflow_id TEXT NOT NULL, + source_revision INTEGER NOT NULL, + source_content_hash TEXT NOT NULL, + source_graph_hash TEXT NOT NULL, + local_conflict_state TEXT NOT NULL CHECK (local_conflict_state IN ('none','identical','id_conflict')), + local_workflow_id TEXT, + local_content_hash TEXT, + local_graph_hash TEXT, + created_by TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'ready' CHECK (status IN ('ready','consumed','expired')), + created_at DATETIME NOT NULL, + expires_at DATETIME NOT NULL, + consumed_at DATETIME +); + +CREATE TABLE workflow_runs ( + id TEXT PRIMARY KEY, + workflow_id TEXT NOT NULL, + workflow_version INTEGER NOT NULL DEFAULT 1, + conversation_id TEXT, + project_id TEXT, + role_id TEXT, + status TEXT NOT NULL, + input_json TEXT, + output_json TEXT, + error TEXT, + pending_hitl_node_id TEXT, + pending_hitl_json TEXT, + started_at DATETIME NOT NULL, + finished_at DATETIME, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE SET NULL +); + +CREATE INDEX idx_assets_domain ON assets(domain); + +CREATE INDEX idx_assets_ip ON assets(ip); + +CREATE INDEX idx_assets_last_scan ON assets(last_scan_at); + +CREATE INDEX idx_assets_last_seen ON assets(last_seen_at); + +CREATE INDEX idx_assets_owner ON assets(owner_user_id); + +CREATE INDEX idx_assets_project ON assets(project_id); + +CREATE INDEX idx_assets_risk_level ON assets(risk_level); + +CREATE INDEX idx_assets_risk_score ON assets(risk_score); + +CREATE INDEX idx_assets_status ON assets(status); + +CREATE INDEX idx_assets_vulnerability_count ON assets(vulnerability_count); + +CREATE INDEX idx_audit_logs_action ON audit_logs(action); + +CREATE INDEX idx_audit_logs_category ON audit_logs(category); + +CREATE INDEX idx_audit_logs_created_at ON audit_logs(created_at); + +CREATE INDEX idx_audit_logs_result ON audit_logs(result); + +CREATE INDEX idx_batch_task_queues_created_at ON batch_task_queues(created_at); + +CREATE INDEX idx_batch_task_queues_title ON batch_task_queues(title); + +CREATE INDEX idx_batch_tasks_queue_id ON batch_tasks(queue_id); + +CREATE INDEX idx_c2_events_category ON c2_events(category); + +CREATE INDEX idx_c2_events_created_at ON c2_events(created_at); + +CREATE INDEX idx_c2_events_session ON c2_events(session_id); + +CREATE INDEX idx_c2_files_session ON c2_files(session_id); + +CREATE INDEX idx_c2_listeners_created_at ON c2_listeners(created_at); + +CREATE INDEX idx_c2_listeners_project_id ON c2_listeners(project_id); + +CREATE INDEX idx_c2_listeners_status ON c2_listeners(status); + +CREATE INDEX idx_c2_payload_artifacts_listener ON c2_payload_artifacts(listener_id); + +CREATE INDEX idx_c2_sessions_last_check_in ON c2_sessions(last_check_in); + +CREATE INDEX idx_c2_sessions_listener ON c2_sessions(listener_id); + +CREATE INDEX idx_c2_sessions_status ON c2_sessions(status); + +CREATE INDEX idx_c2_tasks_conversation ON c2_tasks(conversation_id); + +CREATE INDEX idx_c2_tasks_created_at ON c2_tasks(created_at); + +CREATE INDEX idx_c2_tasks_session ON c2_tasks(session_id); + +CREATE INDEX idx_c2_tasks_status ON c2_tasks(status); + +CREATE INDEX idx_chain_edges_conversation ON attack_chain_edges(conversation_id); + +CREATE INDEX idx_chain_edges_source ON attack_chain_edges(source_node_id); + +CREATE INDEX idx_chain_edges_target ON attack_chain_edges(target_node_id); + +CREATE INDEX idx_chain_nodes_conversation ON attack_chain_nodes(conversation_id); + +CREATE INDEX idx_chat_upload_artifacts_conversation ON chat_upload_artifacts(conversation_id); + +CREATE INDEX idx_chat_upload_artifacts_owner ON chat_upload_artifacts(owner_user_id); + +CREATE INDEX idx_conversations_pinned ON conversations(pinned); + +CREATE INDEX idx_conversations_project_id ON conversations(project_id); + +CREATE INDEX idx_conversations_updated_at ON conversations(updated_at); + +CREATE INDEX idx_knowledge_retrieval_logs_conversation ON knowledge_retrieval_logs(conversation_id); + +CREATE INDEX idx_knowledge_retrieval_logs_created_at ON knowledge_retrieval_logs(created_at); + +CREATE INDEX idx_knowledge_retrieval_logs_message ON knowledge_retrieval_logs(message_id); + +CREATE INDEX idx_messages_conversation_id ON messages(conversation_id); + +CREATE INDEX idx_model_token_usage_conversation ON model_token_usage(conversation_id); + +CREATE INDEX idx_model_token_usage_created_at ON model_token_usage(created_at); + +CREATE INDEX idx_model_token_usage_model ON model_token_usage(model); + +CREATE INDEX idx_model_token_usage_project ON model_token_usage(project_id); + +CREATE INDEX idx_notification_reads_user_read_at ON notification_reads_by_user(user_id, read_at DESC); + +CREATE INDEX idx_process_details_conversation_id ON process_details(conversation_id); + +CREATE INDEX idx_process_details_message_id ON process_details(message_id); + +CREATE INDEX idx_project_fact_edges_project ON project_fact_edges(project_id); + +CREATE INDEX idx_project_fact_edges_source ON project_fact_edges(project_id, source_fact_key); + +CREATE INDEX idx_project_fact_edges_target ON project_fact_edges(project_id, target_fact_key); + +CREATE INDEX idx_project_facts_confidence ON project_facts(confidence); + +CREATE INDEX idx_project_facts_project_id ON project_facts(project_id); + +CREATE INDEX idx_project_facts_related_vuln ON project_facts(related_vulnerability_id); + +CREATE INDEX idx_projects_status ON projects(status); + +CREATE INDEX idx_projects_updated_at ON projects(updated_at); + +CREATE INDEX idx_rbac_assignments_resource ON rbac_resource_assignments(resource_type, resource_id); + +CREATE INDEX idx_rbac_assignments_user_resource ON rbac_resource_assignments(user_id, resource_type, resource_id); + +CREATE INDEX idx_rbac_role_permissions_role ON rbac_role_permissions(role_id); + +CREATE INDEX idx_rbac_user_roles_user ON rbac_user_roles(user_id); + +CREATE INDEX idx_robot_binding_codes_expiry ON robot_binding_codes(expires_at); + +CREATE INDEX idx_robot_user_bindings_user ON robot_user_bindings(rbac_user_id); + +CREATE INDEX idx_robot_user_sessions_updated_at ON robot_user_sessions(updated_at); + +CREATE INDEX idx_tool_executions_conversation ON tool_executions(conversation_id); + +CREATE INDEX idx_tool_executions_owner ON tool_executions(owner_user_id); + +CREATE INDEX idx_tool_executions_start_time ON tool_executions(start_time); + +CREATE INDEX idx_tool_executions_status ON tool_executions(status); + +CREATE INDEX idx_tool_executions_tool_name ON tool_executions(tool_name); + +CREATE INDEX idx_vulnerabilities_conversation_id ON vulnerabilities(conversation_id); + +CREATE INDEX idx_vulnerabilities_conversation_tag ON vulnerabilities(conversation_tag); + +CREATE INDEX idx_vulnerabilities_created_at ON vulnerabilities(created_at); + +CREATE INDEX idx_vulnerabilities_project_id ON vulnerabilities(project_id); + +CREATE INDEX idx_vulnerabilities_severity ON vulnerabilities(severity); + +CREATE INDEX idx_vulnerabilities_status ON vulnerabilities(status); + +CREATE INDEX idx_vulnerabilities_task_tag ON vulnerabilities(task_tag); + +CREATE INDEX idx_webshell_connection_states_updated_at ON webshell_connection_states(updated_at); + +CREATE INDEX idx_webshell_connections_created_at ON webshell_connections(created_at); + +CREATE INDEX idx_webshell_connections_project_id ON webshell_connections(project_id); + +CREATE INDEX idx_workflow_definitions_enabled ON workflow_definitions(enabled); + +CREATE INDEX idx_workflow_definitions_updated_at ON workflow_definitions(updated_at); + +CREATE INDEX idx_workflow_node_runs_run ON workflow_node_runs(run_id); + +CREATE INDEX idx_workflow_package_inspections_creator_expiry ON workflow_package_inspections(created_by, expires_at); + +CREATE INDEX idx_workflow_runs_conversation ON workflow_runs(conversation_id); + +CREATE INDEX idx_workflow_runs_status ON workflow_runs(status); + +CREATE INDEX idx_workflow_runs_workflow ON workflow_runs(workflow_id); + +CREATE UNIQUE INDEX uq_workflow_package_imports_actor_key ON workflow_package_imports(actor_user_id, idempotency_key); + +CREATE UNIQUE INDEX uq_workflow_package_imports_inspection_success ON workflow_package_imports(inspection_id) WHERE result IN ('created','overwritten','renamed','kept_existing','skipped_identical'); diff --git a/internal/database/schemas/knowledge.sql b/internal/database/schemas/knowledge.sql new file mode 100644 index 00000000..11ca413b --- /dev/null +++ b/internal/database/schemas/knowledge.sql @@ -0,0 +1,44 @@ +-- Latest SQLite structure. Update these definitions when adding tables or columns. + +CREATE TABLE knowledge_base_items ( + id TEXT PRIMARY KEY, + category TEXT NOT NULL, + title TEXT NOT NULL, + file_path TEXT NOT NULL, + content TEXT, + created_at DATETIME NOT NULL, + updated_at DATETIME NOT NULL +); + +CREATE TABLE knowledge_embeddings ( + id TEXT PRIMARY KEY, + item_id TEXT NOT NULL, + chunk_index INTEGER NOT NULL, + chunk_text TEXT NOT NULL, + embedding TEXT NOT NULL, + sub_indexes TEXT NOT NULL DEFAULT '', + embedding_model TEXT NOT NULL DEFAULT '', + embedding_dim INTEGER NOT NULL DEFAULT 0, + created_at DATETIME NOT NULL, + FOREIGN KEY (item_id) REFERENCES knowledge_base_items(id) ON DELETE CASCADE +); + +CREATE TABLE knowledge_retrieval_logs ( + id TEXT PRIMARY KEY, + conversation_id TEXT, + message_id TEXT, + query TEXT NOT NULL, + risk_type TEXT, + retrieved_items TEXT, + created_at DATETIME NOT NULL +); + +CREATE INDEX idx_knowledge_embeddings_item_id ON knowledge_embeddings(item_id); + +CREATE INDEX idx_knowledge_items_category ON knowledge_base_items(category); + +CREATE INDEX idx_knowledge_retrieval_logs_conversation ON knowledge_retrieval_logs(conversation_id); + +CREATE INDEX idx_knowledge_retrieval_logs_created_at ON knowledge_retrieval_logs(created_at); + +CREATE INDEX idx_knowledge_retrieval_logs_message ON knowledge_retrieval_logs(message_id); diff --git a/internal/database/workflow.go b/internal/database/workflow.go index 5958f4ba..33ab9d0b 100644 --- a/internal/database/workflow.go +++ b/internal/database/workflow.go @@ -466,27 +466,6 @@ func (db *DB) ListWorkflowRunsAwaitingHITLFiltered(conversationID string, limit return out, rows.Err() } -func (db *DB) migrateWorkflowRunsTable() error { - cols := []struct{ name, ddl string }{ - {"pending_hitl_node_id", "ALTER TABLE workflow_runs ADD COLUMN pending_hitl_node_id TEXT"}, - {"pending_hitl_json", "ALTER TABLE workflow_runs ADD COLUMN pending_hitl_json TEXT"}, - } - for _, col := range cols { - var count int - err := db.QueryRow("SELECT COUNT(*) FROM pragma_table_info('workflow_runs') WHERE name=?", col.name).Scan(&count) - if err != nil || count > 0 { - continue - } - if _, err := db.Exec(col.ddl); err != nil { - errMsg := strings.ToLower(err.Error()) - if !strings.Contains(errMsg, "duplicate column") && !strings.Contains(errMsg, "already exists") { - return err - } - } - } - return nil -} - func nullString(v string) interface{} { v = strings.TrimSpace(v) if v == "" { diff --git a/internal/handler/hitl.go b/internal/handler/hitl.go index 44b8b4a4..143ab47a 100644 --- a/internal/handler/hitl.go +++ b/internal/handler/hitl.go @@ -63,38 +63,11 @@ func NewHITLManager(db *database.DB, logger *zap.Logger) *HITLManager { } } +// EnsureSchema performs historical recovery after NewDB has checked all tables. func (m *HITLManager) EnsureSchema() error { - if _, err := m.db.Exec(` -CREATE TABLE IF NOT EXISTS hitl_interrupts ( - id TEXT PRIMARY KEY, - conversation_id TEXT NOT NULL, - message_id TEXT, - mode TEXT NOT NULL, - tool_name TEXT NOT NULL, - tool_call_id TEXT, - payload TEXT, - status TEXT NOT NULL, - reviewer TEXT NOT NULL DEFAULT 'human', - decision TEXT, - decision_comment TEXT, - created_at DATETIME NOT NULL, - decided_at DATETIME -);`); err != nil { + if err := m.backfillHitlReviewers(); err != nil { return err } - _, err := m.db.Exec(` -CREATE TABLE IF NOT EXISTS hitl_conversation_configs ( - conversation_id TEXT PRIMARY KEY, - enabled INTEGER NOT NULL DEFAULT 0, - mode TEXT NOT NULL DEFAULT 'off', - sensitive_tools TEXT NOT NULL DEFAULT '[]', - timeout_seconds INTEGER NOT NULL DEFAULT 0, - updated_at DATETIME NOT NULL -);`) - if err != nil { - return err - } - m.migrateHitlSchemaColumns() // On startup, cancel all orphaned pending interrupts from previous process. // Their in-memory channels are gone, so they can never be resolved. diff --git a/internal/handler/hitl_logs.go b/internal/handler/hitl_logs.go index 773e35dc..1448e80d 100644 --- a/internal/handler/hitl_logs.go +++ b/internal/handler/hitl_logs.go @@ -37,12 +37,10 @@ func normalizeHitlDecidedBy(v string) string { } } -func (m *HITLManager) migrateHitlSchemaColumns() { - _, _ = m.db.Exec(`ALTER TABLE hitl_interrupts ADD COLUMN decided_by TEXT NOT NULL DEFAULT 'human'`) - _, _ = m.db.Exec(`ALTER TABLE hitl_interrupts ADD COLUMN reviewer TEXT NOT NULL DEFAULT 'human'`) - _, _ = m.db.Exec(`UPDATE hitl_interrupts SET reviewer='audit_agent' - WHERE COALESCE(decided_by, '') IN ('audit_agent', 'agent', 'ai')`) - _, _ = m.db.Exec(`ALTER TABLE hitl_conversation_configs ADD COLUMN reviewer TEXT NOT NULL DEFAULT 'human'`) +func (m *HITLManager) backfillHitlReviewers() error { + _, err := m.db.Exec(`UPDATE hitl_interrupts SET reviewer='audit_agent' + WHERE COALESCE(decided_by, '') IN ('audit_agent', 'agent', 'ai')`) + return err } func hitlInterruptRowToMap( diff --git a/internal/handler/notification.go b/internal/handler/notification.go index 21f99d9b..278c7836 100644 --- a/internal/handler/notification.go +++ b/internal/handler/notification.go @@ -582,25 +582,6 @@ func countUnread(items []NotificationSummaryItem) int { return total } -func createNotificationReadTableIfNeeded(db *database.DB) error { - if db == nil { - return fmt.Errorf("db is nil") - } - _, err := db.Exec(` - CREATE TABLE IF NOT EXISTS notification_reads_by_user ( - user_id TEXT NOT NULL, - event_id TEXT NOT NULL, - read_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - PRIMARY KEY(user_id, event_id) - ); - `) - if err != nil { - return err - } - _, idxErr := db.Exec(`CREATE INDEX IF NOT EXISTS idx_notification_reads_user_read_at ON notification_reads_by_user(user_id, read_at DESC);`) - return idxErr -} - func pruneNotificationReads(db *database.DB, userID string, maxRows int) error { if db == nil { return fmt.Errorf("db is nil") @@ -648,7 +629,7 @@ func normalizeMarkableEventID(id string) (string, bool) { // MarkRead 按事件 ID 标记已读 func (h *NotificationHandler) MarkRead(c *gin.Context) { - if err := createNotificationReadTableIfNeeded(h.db); err != nil { + if h.db == nil { c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to prepare notification read table"}) return } @@ -713,12 +694,6 @@ func (h *NotificationHandler) GetSummary(c *gin.Context) { return } - if err := createNotificationReadTableIfNeeded(h.db); err != nil { - h.logger.Warn("初始化通知已读表失败", zap.Error(err)) - c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to initialize notification read table"}) - return - } - english := requestWantsEnglish(c) sinceMs := normalizeSinceMs(parseSinceMs(c.Query("since"))) limit, _ := strconv.Atoi(strings.TrimSpace(c.DefaultQuery("limit", "50"))) diff --git a/internal/knowledge/schema_migrate.go b/internal/knowledge/schema_migrate.go index 85fd26e2..371dc9f2 100644 --- a/internal/knowledge/schema_migrate.go +++ b/internal/knowledge/schema_migrate.go @@ -1,51 +1,17 @@ package knowledge import ( + "cyberstrike-ai/internal/database" "database/sql" - "fmt" ) -// EnsureKnowledgeEmbeddingsSchema migrates knowledge_embeddings for sub_indexes + embedding metadata. +// EnsureKnowledgeEmbeddingsSchema checks against the canonical SQLite schema; +// callers no longer maintain an independent list of embedding columns. func EnsureKnowledgeEmbeddingsSchema(db *sql.DB) error { - if db == nil { - return fmt.Errorf("db is nil") - } - var n int - if err := db.QueryRow(`SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='knowledge_embeddings'`).Scan(&n); err != nil { - return err - } - if n == 0 { - return nil - } - if err := addKnowledgeEmbeddingsColumnIfMissing(db, "sub_indexes", - `ALTER TABLE knowledge_embeddings ADD COLUMN sub_indexes TEXT NOT NULL DEFAULT ''`); err != nil { - return err - } - if err := addKnowledgeEmbeddingsColumnIfMissing(db, "embedding_model", - `ALTER TABLE knowledge_embeddings ADD COLUMN embedding_model TEXT NOT NULL DEFAULT ''`); err != nil { - return err - } - if err := addKnowledgeEmbeddingsColumnIfMissing(db, "embedding_dim", - `ALTER TABLE knowledge_embeddings ADD COLUMN embedding_dim INTEGER NOT NULL DEFAULT 0`); err != nil { - return err - } - return nil + return database.EnsureKnowledgeSchema(db) } -func addKnowledgeEmbeddingsColumnIfMissing(db *sql.DB, column, alterSQL string) error { - var colCount int - q := `SELECT COUNT(*) FROM pragma_table_info('knowledge_embeddings') WHERE name = ?` - if err := db.QueryRow(q, column).Scan(&colCount); err != nil { - return err - } - if colCount > 0 { - return nil - } - _, err := db.Exec(alterSQL) - return err -} - -// ensureKnowledgeEmbeddingsSubIndexesColumn 向后兼容；请使用 [EnsureKnowledgeEmbeddingsSchema]。 +// ensureKnowledgeEmbeddingsSubIndexesColumn is retained for compatibility. func ensureKnowledgeEmbeddingsSubIndexesColumn(db *sql.DB) error { return EnsureKnowledgeEmbeddingsSchema(db) }