Add files via upload

This commit is contained in:
公明
2026-05-26 14:27:44 +08:00
committed by GitHub
parent 44ced98863
commit 1b1aed1699
11 changed files with 790 additions and 36 deletions
+5 -3
View File
@@ -17,6 +17,7 @@ import (
"cyberstrike-ai/internal/config" "cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp" "cyberstrike-ai/internal/mcp"
"cyberstrike-ai/internal/mcp/builtin" "cyberstrike-ai/internal/mcp/builtin"
"cyberstrike-ai/internal/project"
"cyberstrike-ai/internal/openai" "cyberstrike-ai/internal/openai"
"cyberstrike-ai/internal/security" "cyberstrike-ai/internal/security"
"cyberstrike-ai/internal/storage" "cyberstrike-ai/internal/storage"
@@ -365,12 +366,12 @@ type ProgressCallback func(eventType, message string, data interface{})
// AgentLoop 执行Agent循环 // AgentLoop 执行Agent循环
func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages []ChatMessage) (*AgentLoopResult, error) { func (a *Agent) AgentLoop(ctx context.Context, userInput string, historyMessages []ChatMessage) (*AgentLoopResult, error) {
return a.AgentLoopWithProgress(ctx, userInput, historyMessages, "", nil, nil) return a.AgentLoopWithProgress(ctx, userInput, historyMessages, "", nil, nil, "")
} }
// AgentLoopWithConversationID 执行Agent循环(带对话ID // AgentLoopWithConversationID 执行Agent循环(带对话ID
func (a *Agent) AgentLoopWithConversationID(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string) (*AgentLoopResult, error) { func (a *Agent) AgentLoopWithConversationID(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string) (*AgentLoopResult, error) {
return a.AgentLoopWithProgress(ctx, userInput, historyMessages, conversationID, nil, nil) return a.AgentLoopWithProgress(ctx, userInput, historyMessages, conversationID, nil, nil, "")
} }
// EinoSingleAgentSystemInstruction 供 Eino adk.ChatModelAgent.Instruction 使用,与 AgentLoopWithProgress 首条 system 对齐(含 system_prompt_path)。 // EinoSingleAgentSystemInstruction 供 Eino adk.ChatModelAgent.Instruction 使用,与 AgentLoopWithProgress 首条 system 对齐(含 system_prompt_path)。
@@ -396,7 +397,7 @@ func (a *Agent) EinoSingleAgentSystemInstruction() string {
} }
// AgentLoopWithProgress 执行Agent循环(带进度回调和对话ID) // AgentLoopWithProgress 执行Agent循环(带进度回调和对话ID)
func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string, callback ProgressCallback, roleTools []string) (*AgentLoopResult, error) { func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, historyMessages []ChatMessage, conversationID string, callback ProgressCallback, roleTools []string, systemPromptExtra string) (*AgentLoopResult, error) {
ctx = withAgentConversationID(ctx, conversationID) ctx = withAgentConversationID(ctx, conversationID)
// 设置当前对话ID(兼容未走 context 的旧路径;并发会话应以 context 为准) // 设置当前对话ID(兼容未走 context 的旧路径;并发会话应以 context 为准)
a.mu.Lock() a.mu.Lock()
@@ -426,6 +427,7 @@ func (a *Agent) AgentLoopWithProgress(ctx context.Context, userInput string, his
} }
} }
} }
systemPrompt = project.AppendSystemPromptBlock(systemPrompt, systemPromptExtra)
messages := []ChatMessage{ messages := []ChatMessage{
{ {
@@ -105,11 +105,15 @@ func DefaultSingleAgentSystemPrompt() string {
- 若最近一步得到 404/空结果/无效响应,不得直接结束;至少再进行一次“同目标不同策略”的验证(如变更路径、参数、请求方法、上下文来源)。 - 若最近一步得到 404/空结果/无效响应,不得直接结束;至少再进行一次“同目标不同策略”的验证(如变更路径、参数、请求方法、上下文来源)。
- 避免无效空转:同一工具+同类参数连续失败 3 次后,必须切换策略(改工具、改入口、改假设)并说明切换原因。 - 避免无效空转:同一工具+同类参数连续失败 3 次后,必须切换策略(改工具、改入口、改假设)并说明切换原因。
## 漏洞记录 ## 项目黑板(事实)与漏洞记录(分离)
发现有效漏洞时,必须使` + builtin.ToolRecordVulnerability + ` 记录:标题、描述、严重程度、类型、目标、证明(POC)、影响、修复建议。 当前对话若已绑定项目,系统会自动注入「项目黑板索引」(仅 fact_key + 摘要)。**摘要不足时必须` + builtin.ToolGetProjectFact + `(fact_key) 获取 body,禁止凭摘要臆造细节。**
严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。记录后可在授权范围内继续测试 - **环境/目标/认证等认知**(非正式漏洞条目):使用 ` + builtin.ToolUpsertProjectFact + `fact_key 建议 ` + "`category/slug`" + `(如 target/primary_domain),同 key 覆盖更新
- **可交付漏洞**:使用 ` + builtin.ToolRecordVulnerability + `,含标题、严重程度、类型、目标、证明(POC)、影响、修复建议。记前可先 ` + builtin.ToolListVulnerabilities + ` 查重,详情用 ` + builtin.ToolGetVulnerability + `(id)(默认仅当前项目/会话)。
- 同一发现可能需**各记一次**(事实记上下文,漏洞记正式 findings)。误报用 ` + builtin.ToolDeprecateProjectFact + ` 或漏洞状态 false_positive。
严重程度:critical / high / medium / low / info。证明须含足够证据(请求响应、截图、命令输出等)。
## 技能库(Skills)与知识库 ## 技能库(Skills)与知识库
+26
View File
@@ -36,6 +36,32 @@ type Config struct {
SkillsDir string `yaml:"skills_dir,omitempty" json:"skills_dir,omitempty"` // Skills配置文件目录 SkillsDir string `yaml:"skills_dir,omitempty" json:"skills_dir,omitempty"` // Skills配置文件目录
AgentsDir string `yaml:"agents_dir,omitempty" json:"agents_dir,omitempty"` // 多代理子 Agent Markdown 定义目录(*.mdYAML front matter AgentsDir string `yaml:"agents_dir,omitempty" json:"agents_dir,omitempty"` // 多代理子 Agent Markdown 定义目录(*.mdYAML front matter
MultiAgent MultiAgentConfig `yaml:"multi_agent,omitempty" json:"multi_agent,omitempty"` MultiAgent MultiAgentConfig `yaml:"multi_agent,omitempty" json:"multi_agent,omitempty"`
Project ProjectConfig `yaml:"project,omitempty" json:"project,omitempty"`
}
// ProjectConfig 项目黑板(跨对话共享事实)配置。
type ProjectConfig struct {
Enabled bool `yaml:"enabled" json:"enabled"`
DefaultProjectID string `yaml:"default_project_id,omitempty" json:"default_project_id,omitempty"` // 机器人/批量等无显式项目时绑定的默认项目
FactIndexMaxRunes int `yaml:"fact_index_max_runes,omitempty" json:"fact_index_max_runes,omitempty"`
FactSummaryMaxRunes int `yaml:"fact_summary_max_runes,omitempty" json:"fact_summary_max_runes,omitempty"`
DefaultInjectDeprecated bool `yaml:"default_inject_deprecated,omitempty" json:"default_inject_deprecated,omitempty"`
}
// FactIndexMaxRunesEffective 自动注入黑板索引的最大 rune 数。
func (c ProjectConfig) FactIndexMaxRunesEffective() int {
if c.FactIndexMaxRunes <= 0 {
return 3500
}
return c.FactIndexMaxRunes
}
// FactSummaryMaxRunesEffective upsert 时 summary 最大 rune 数。
func (c ProjectConfig) FactSummaryMaxRunesEffective() int {
if c.FactSummaryMaxRunes <= 0 {
return 120
}
return c.FactSummaryMaxRunes
} }
// MultiAgentConfig 基于 CloudWeGo Eino adk/prebuilt 的多代理编排(deep | plan_execute | supervisor,与单 Agent /agent-loop 并存)。 // MultiAgentConfig 基于 CloudWeGo Eino adk/prebuilt 的多代理编排(deep | plan_execute | supervisor,与单 Agent /agent-loop 并存)。
+14 -8
View File
@@ -22,6 +22,7 @@ type BatchTaskQueueRow struct {
LastScheduleTriggerAt sql.NullTime LastScheduleTriggerAt sql.NullTime
LastScheduleError sql.NullString LastScheduleError sql.NullString
LastRunError sql.NullString LastRunError sql.NullString
ProjectID sql.NullString
Status string Status string
CreatedAt time.Time CreatedAt time.Time
StartedAt sql.NullTime StartedAt sql.NullTime
@@ -51,6 +52,7 @@ func (db *DB) CreateBatchQueue(
scheduleMode string, scheduleMode string,
cronExpr string, cronExpr string,
nextRunAt *time.Time, nextRunAt *time.Time,
projectID string,
tasks []map[string]interface{}, tasks []map[string]interface{},
) error { ) error {
tx, err := db.Begin() tx, err := db.Begin()
@@ -65,9 +67,13 @@ func (db *DB) CreateBatchQueue(
nextRunAtValue = *nextRunAt nextRunAtValue = *nextRunAt
} }
var projectIDVal interface{}
if strings.TrimSpace(projectID) != "" {
projectIDVal = strings.TrimSpace(projectID)
}
_, err = tx.Exec( _, err = tx.Exec(
"INSERT INTO batch_task_queues (id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, status, created_at, current_index) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", "INSERT INTO batch_task_queues (id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, project_id, status, created_at, current_index) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
queueID, title, role, agentMode, scheduleMode, cronExpr, nextRunAtValue, 1, "pending", now, 0, queueID, title, role, agentMode, scheduleMode, cronExpr, nextRunAtValue, 1, projectIDVal, "pending", now, 0,
) )
if err != nil { if err != nil {
return fmt.Errorf("创建批量任务队列失败: %w", err) return fmt.Errorf("创建批量任务队列失败: %w", err)
@@ -101,9 +107,9 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
var row BatchTaskQueueRow var row BatchTaskQueueRow
var createdAt string var createdAt string
err := db.QueryRow( err := db.QueryRow(
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE id = ?", "SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE id = ?",
queueID, queueID,
).Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex) ).Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex)
if err == sql.ErrNoRows { if err == sql.ErrNoRows {
return nil, nil return nil, nil
} }
@@ -127,7 +133,7 @@ func (db *DB) GetBatchQueue(queueID string) (*BatchTaskQueueRow, error) {
// GetAllBatchQueues 获取所有批量任务队列 // GetAllBatchQueues 获取所有批量任务队列
func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) { func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
rows, err := db.Query( rows, err := db.Query(
"SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, status, created_at, started_at, completed_at, current_index FROM batch_task_queues ORDER BY created_at DESC", "SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues ORDER BY created_at DESC",
) )
if err != nil { if err != nil {
return nil, fmt.Errorf("查询批量任务队列列表失败: %w", err) return nil, fmt.Errorf("查询批量任务队列列表失败: %w", err)
@@ -138,7 +144,7 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
for rows.Next() { for rows.Next() {
var row BatchTaskQueueRow var row BatchTaskQueueRow
var createdAt string var createdAt string
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil { if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
return nil, fmt.Errorf("扫描批量任务队列失败: %w", err) return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
} }
parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt) parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
@@ -158,7 +164,7 @@ func (db *DB) GetAllBatchQueues() ([]*BatchTaskQueueRow, error) {
// ListBatchQueues 列出批量任务队列(支持筛选和分页) // ListBatchQueues 列出批量任务队列(支持筛选和分页)
func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueueRow, error) { func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*BatchTaskQueueRow, error) {
query := "SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE 1=1" query := "SELECT id, title, role, agent_mode, schedule_mode, cron_expr, next_run_at, schedule_enabled, last_schedule_trigger_at, last_schedule_error, last_run_error, project_id, status, created_at, started_at, completed_at, current_index FROM batch_task_queues WHERE 1=1"
args := []interface{}{} args := []interface{}{}
// 状态筛选 // 状态筛选
@@ -186,7 +192,7 @@ func (db *DB) ListBatchQueues(limit, offset int, status, keyword string) ([]*Bat
for rows.Next() { for rows.Next() {
var row BatchTaskQueueRow var row BatchTaskQueueRow
var createdAt string var createdAt string
if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil { if err := rows.Scan(&row.ID, &row.Title, &row.Role, &row.AgentMode, &row.ScheduleMode, &row.CronExpr, &row.NextRunAt, &row.ScheduleEnabled, &row.LastScheduleTriggerAt, &row.LastScheduleError, &row.LastRunError, &row.ProjectID, &row.Status, &createdAt, &row.StartedAt, &row.CompletedAt, &row.CurrentIndex); err != nil {
return nil, fmt.Errorf("扫描批量任务队列失败: %w", err) return nil, fmt.Errorf("扫描批量任务队列失败: %w", err)
} }
parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt) parsedTime, parseErr := time.Parse("2006-01-02 15:04:05", createdAt)
+45 -12
View File
@@ -17,6 +17,7 @@ import (
type Conversation struct { type Conversation struct {
ID string `json:"id"` ID string `json:"id"`
Title string `json:"title"` Title string `json:"title"`
ProjectID string `json:"projectId,omitempty"`
Pinned bool `json:"pinned"` Pinned bool `json:"pinned"`
CreatedAt time.Time `json:"createdAt"` CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"` UpdatedAt time.Time `json:"updatedAt"`
@@ -46,13 +47,32 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string,
id := uuid.New().String() id := uuid.New().String()
now := time.Now() now := time.Now()
projectID := strings.TrimSpace(meta.ProjectID)
if projectID != "" {
if _, err := db.GetProject(projectID); err != nil {
return nil, err
}
}
var err error var err error
if webshellConnectionID != "" { wsID := strings.TrimSpace(webshellConnectionID)
switch {
case wsID != "" && projectID != "":
_, err = db.Exec(
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id, project_id) VALUES (?, ?, ?, ?, ?, ?)",
id, title, now, now, wsID, projectID,
)
case wsID != "":
_, err = db.Exec( _, err = db.Exec(
"INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id) VALUES (?, ?, ?, ?, ?)", "INSERT INTO conversations (id, title, created_at, updated_at, webshell_connection_id) VALUES (?, ?, ?, ?, ?)",
id, title, now, now, webshellConnectionID, id, title, now, now, wsID,
) )
} else { case projectID != "":
_, err = db.Exec(
"INSERT INTO conversations (id, title, created_at, updated_at, project_id) VALUES (?, ?, ?, ?, ?)",
id, title, now, now, projectID,
)
default:
_, err = db.Exec( _, err = db.Exec(
"INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)", "INSERT INTO conversations (id, title, created_at, updated_at) VALUES (?, ?, ?, ?)",
id, title, now, now, id, title, now, now,
@@ -65,11 +85,12 @@ func (db *DB) CreateConversationWithWebshell(webshellConnectionID, title string,
conv := &Conversation{ conv := &Conversation{
ID: id, ID: id,
Title: title, Title: title,
ProjectID: projectID,
CreatedAt: now, CreatedAt: now,
UpdatedAt: now, UpdatedAt: now,
} }
if webshellConnectionID != "" { if wsID != "" {
meta.WebShellConnectionID = webshellConnectionID meta.WebShellConnectionID = wsID
} }
notifyConversationCreated(conv, meta) notifyConversationCreated(conv, meta)
return conv, nil return conv, nil
@@ -210,16 +231,20 @@ func (db *DB) GetConversation(id string) (*Conversation, error) {
var createdAt, updatedAt string var createdAt, updatedAt string
var pinned int var pinned int
var projectID sql.NullString
err := db.QueryRow( err := db.QueryRow(
"SELECT id, title, pinned, created_at, updated_at FROM conversations WHERE id = ?", "SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
id, id,
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt) ).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
if err != nil { if err != nil {
if err == sql.ErrNoRows { if err == sql.ErrNoRows {
return nil, fmt.Errorf("对话不存在") return nil, fmt.Errorf("对话不存在")
} }
return nil, fmt.Errorf("查询对话失败: %w", err) return nil, fmt.Errorf("查询对话失败: %w", err)
} }
if projectID.Valid {
conv.ProjectID = strings.TrimSpace(projectID.String)
}
// 尝试多种时间格式解析 // 尝试多种时间格式解析
var err1, err2 error var err1, err2 error
@@ -292,16 +317,20 @@ func (db *DB) GetConversationLite(id string) (*Conversation, error) {
var createdAt, updatedAt string var createdAt, updatedAt string
var pinned int var pinned int
var projectID sql.NullString
err := db.QueryRow( err := db.QueryRow(
"SELECT id, title, pinned, created_at, updated_at FROM conversations WHERE id = ?", "SELECT id, title, pinned, created_at, updated_at, project_id FROM conversations WHERE id = ?",
id, id,
).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt) ).Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID)
if err != nil { if err != nil {
if err == sql.ErrNoRows { if err == sql.ErrNoRows {
return nil, fmt.Errorf("对话不存在") return nil, fmt.Errorf("对话不存在")
} }
return nil, fmt.Errorf("查询对话失败: %w", err) return nil, fmt.Errorf("查询对话失败: %w", err)
} }
if projectID.Valid {
conv.ProjectID = strings.TrimSpace(projectID.String)
}
// 尝试多种时间格式解析 // 尝试多种时间格式解析
var err1, err2 error var err1, err2 error
@@ -341,7 +370,7 @@ func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversati
// 使用 EXISTS 子查询代替 LEFT JOIN + DISTINCT,避免大表笛卡尔积 // 使用 EXISTS 子查询代替 LEFT JOIN + DISTINCT,避免大表笛卡尔积
searchPattern := "%" + search + "%" searchPattern := "%" + search + "%"
rows, err = db.Query( rows, err = db.Query(
`SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at `SELECT c.id, c.title, COALESCE(c.pinned, 0), c.created_at, c.updated_at, c.project_id
FROM conversations c FROM conversations c
WHERE c.title LIKE ? WHERE c.title LIKE ?
OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?) OR EXISTS (SELECT 1 FROM messages m WHERE m.conversation_id = c.id AND m.content LIKE ?)
@@ -351,7 +380,7 @@ func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversati
) )
} else { } else {
rows, err = db.Query( rows, err = db.Query(
"SELECT id, title, COALESCE(pinned, 0), created_at, updated_at FROM conversations ORDER BY updated_at DESC LIMIT ? OFFSET ?", "SELECT id, title, COALESCE(pinned, 0), created_at, updated_at, project_id FROM conversations ORDER BY updated_at DESC LIMIT ? OFFSET ?",
limit, offset, limit, offset,
) )
} }
@@ -366,10 +395,14 @@ func (db *DB) ListConversations(limit, offset int, search string) ([]*Conversati
var conv Conversation var conv Conversation
var createdAt, updatedAt string var createdAt, updatedAt string
var pinned int var pinned int
var projectID sql.NullString
if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt); err != nil { if err := rows.Scan(&conv.ID, &conv.Title, &pinned, &createdAt, &updatedAt, &projectID); err != nil {
return nil, fmt.Errorf("扫描对话失败: %w", err) return nil, fmt.Errorf("扫描对话失败: %w", err)
} }
if projectID.Valid {
conv.ProjectID = strings.TrimSpace(projectID.String)
}
// 尝试多种时间格式解析 // 尝试多种时间格式解析
var err1, err2 error var err1, err2 error
@@ -4,6 +4,7 @@ package database
type ConversationCreateMeta struct { type ConversationCreateMeta struct {
Source string Source string
WebShellConnectionID string WebShellConnectionID string
ProjectID string
ClientIP string ClientIP string
SessionHint string SessionHint string
} }
+98
View File
@@ -213,6 +213,40 @@ func (db *DB) initTables() error {
FOREIGN KEY (conversation_id) REFERENCES conversations(id) ON DELETE CASCADE 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,
supersedes_fact_id TEXT,
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)
);`
// 创建漏洞表 // 创建漏洞表
createVulnerabilitiesTable := ` createVulnerabilitiesTable := `
CREATE TABLE IF NOT EXISTS vulnerabilities ( CREATE TABLE IF NOT EXISTS vulnerabilities (
@@ -445,6 +479,12 @@ func (db *DB) initTables() error {
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_severity ON vulnerabilities(severity); 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_status ON vulnerabilities(status);
CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at); CREATE INDEX IF NOT EXISTS idx_vulnerabilities_created_at ON vulnerabilities(created_at);
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_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_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_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_batch_task_queues_title ON batch_task_queues(title);
@@ -516,6 +556,14 @@ func (db *DB) initTables() error {
return fmt.Errorf("创建robot_user_sessions表失败: %w", err) 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(createVulnerabilitiesTable); err != nil { if _, err := db.Exec(createVulnerabilitiesTable); err != nil {
return fmt.Errorf("创建vulnerabilities表失败: %w", err) return fmt.Errorf("创建vulnerabilities表失败: %w", err)
} }
@@ -583,6 +631,10 @@ func (db *DB) initTables() error {
// 不返回错误,允许继续运行 // 不返回错误,允许继续运行
} }
if err := db.migrateProjectsTable(); err != nil {
db.logger.Warn("迁移projects相关表失败", zap.Error(err))
}
if err := db.migrateWebshellConnectionsTable(); err != nil { if err := db.migrateWebshellConnectionsTable(); err != nil {
db.logger.Warn("迁移webshell_connections表失败", zap.Error(err)) db.logger.Warn("迁移webshell_connections表失败", zap.Error(err))
// 不返回错误,允许继续运行 // 不返回错误,允许继续运行
@@ -930,6 +982,51 @@ func (db *DB) migrateBatchTaskQueuesTable() error {
} }
} }
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))
}
}
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 return nil
} }
@@ -941,6 +1038,7 @@ func (db *DB) migrateVulnerabilitiesTable() error {
}{ }{
{name: "conversation_tag", stmt: "ALTER TABLE vulnerabilities ADD COLUMN conversation_tag TEXT"}, {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: "task_tag", stmt: "ALTER TABLE vulnerabilities ADD COLUMN task_tag TEXT"},
{name: "project_id", stmt: "ALTER TABLE vulnerabilities ADD COLUMN project_id TEXT"},
} }
for _, col := range columns { for _, col := range columns {
+451
View File
@@ -0,0 +1,451 @@
package database
import (
"database/sql"
"fmt"
"regexp"
"strings"
"time"
"github.com/google/uuid"
)
var factKeyPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9._/-]*$`)
// ValidateFactKey 校验事实 key(项目内唯一标识)。
func ValidateFactKey(key string) error {
key = strings.TrimSpace(key)
if key == "" {
return fmt.Errorf("fact_key 不能为空")
}
if len(key) > 128 {
return fmt.Errorf("fact_key 过长(最多 128 字符)")
}
if !factKeyPattern.MatchString(key) {
return fmt.Errorf("fact_key 格式无效,仅允许小写字母、数字及 . _ / -,且须以小写字母或数字开头")
}
return nil
}
// Project 渗透测试项目(跨对话共享黑板)。
type Project struct {
ID string `json:"id"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
ScopeJSON string `json:"scope_json,omitempty"`
Status string `json:"status"` // active | archived
Pinned bool `json:"pinned"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// ProjectFact 项目事实(黑板条目)。
type ProjectFact struct {
ID string `json:"id"`
ProjectID string `json:"project_id"`
FactKey string `json:"fact_key"`
Category string `json:"category"`
Summary string `json:"summary"`
Body string `json:"body"`
Confidence string `json:"confidence"` // confirmed | tentative | deprecated
SourceConversationID string `json:"source_conversation_id,omitempty"`
SourceMessageID string `json:"source_message_id,omitempty"`
Pinned bool `json:"pinned"`
SupersedesFactID string `json:"supersedes_fact_id,omitempty"`
RelatedVulnerabilityID string `json:"related_vulnerability_id,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// ProjectFactListFilter 事实列表筛选。
type ProjectFactListFilter struct {
Category string
Confidence string
Search string
}
// CreateProject 创建项目。
func (db *DB) CreateProject(p *Project) (*Project, error) {
if p.ID == "" {
p.ID = uuid.New().String()
}
if strings.TrimSpace(p.Status) == "" {
p.Status = "active"
}
now := time.Now()
if p.CreatedAt.IsZero() {
p.CreatedAt = now
}
p.UpdatedAt = now
_, err := db.Exec(
`INSERT INTO projects (id, name, description, scope_json, status, pinned, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)`,
p.ID, p.Name, p.Description, p.ScopeJSON, p.Status, boolToInt(p.Pinned), p.CreatedAt, p.UpdatedAt,
)
if err != nil {
return nil, fmt.Errorf("创建项目失败: %w", err)
}
return p, nil
}
// GetProject 获取项目。
func (db *DB) GetProject(id string) (*Project, error) {
var p Project
var pinned int
var createdAt, updatedAt string
err := db.QueryRow(
`SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
FROM projects WHERE id = ?`, id,
).Scan(&p.ID, &p.Name, &p.Description, &p.ScopeJSON, &p.Status, &pinned, &createdAt, &updatedAt)
if err != nil {
if err == sql.ErrNoRows {
return nil, fmt.Errorf("项目不存在")
}
return nil, fmt.Errorf("获取项目失败: %w", err)
}
p.Pinned = pinned != 0
p.CreatedAt = parseDBTime(createdAt)
p.UpdatedAt = parseDBTime(updatedAt)
return &p, nil
}
// ListProjects 列出项目。
func (db *DB) ListProjects(status string, limit, offset int) ([]*Project, error) {
if limit <= 0 {
limit = 200
}
query := `SELECT id, name, COALESCE(description,''), COALESCE(scope_json,''), status, pinned, created_at, updated_at
FROM projects WHERE 1=1`
args := []interface{}{}
if s := strings.TrimSpace(status); s != "" {
query += " AND status = ?"
args = append(args, s)
}
query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?"
args = append(args, limit, offset)
rows, err := db.Query(query, args...)
if err != nil {
return nil, fmt.Errorf("列出项目失败: %w", err)
}
defer rows.Close()
var out []*Project
for rows.Next() {
var p Project
var pinned int
var createdAt, updatedAt string
if err := rows.Scan(&p.ID, &p.Name, &p.Description, &p.ScopeJSON, &p.Status, &pinned, &createdAt, &updatedAt); err != nil {
return nil, err
}
p.Pinned = pinned != 0
p.CreatedAt = parseDBTime(createdAt)
p.UpdatedAt = parseDBTime(updatedAt)
out = append(out, &p)
}
return out, rows.Err()
}
// UpdateProject 更新项目。
func (db *DB) UpdateProject(p *Project) error {
p.UpdatedAt = time.Now()
_, err := db.Exec(
`UPDATE projects SET name = ?, description = ?, scope_json = ?, status = ?, pinned = ?, updated_at = ? WHERE id = ?`,
p.Name, p.Description, p.ScopeJSON, p.Status, boolToInt(p.Pinned), p.UpdatedAt, p.ID,
)
if err != nil {
return fmt.Errorf("更新项目失败: %w", err)
}
return nil
}
// DeleteProject 删除项目(级联删除事实;对话 project_id 置空由 FK 处理)。
func (db *DB) DeleteProject(id string) error {
_, err := db.Exec(`DELETE FROM projects WHERE id = ?`, id)
if err != nil {
return fmt.Errorf("删除项目失败: %w", err)
}
return nil
}
// GetConversationProjectID 返回对话绑定的项目 ID。
func (db *DB) GetConversationProjectID(conversationID string) (string, error) {
var pid sql.NullString
err := db.QueryRow(`SELECT project_id FROM conversations WHERE id = ?`, conversationID).Scan(&pid)
if err != nil {
if err == sql.ErrNoRows {
return "", fmt.Errorf("对话不存在")
}
return "", err
}
if pid.Valid {
return strings.TrimSpace(pid.String), nil
}
return "", nil
}
// SetConversationProjectID 设置对话所属项目(空字符串表示解除绑定)。
func (db *DB) SetConversationProjectID(conversationID, projectID string) error {
projectID = strings.TrimSpace(projectID)
if projectID != "" {
if _, err := db.GetProject(projectID); err != nil {
return err
}
}
var val interface{}
if projectID == "" {
val = nil
} else {
val = projectID
}
_, err := db.Exec(`UPDATE conversations SET project_id = ?, updated_at = ? WHERE id = ?`, val, time.Now(), conversationID)
if err != nil {
return fmt.Errorf("设置对话项目失败: %w", err)
}
return nil
}
// ListProjectFactsForIndex 列出用于黑板索引注入的事实(不含 deprecated,除非 includeDeprecated)。
func (db *DB) ListProjectFactsForIndex(projectID string, includeDeprecated bool) ([]*ProjectFact, error) {
query := `SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
FROM project_facts WHERE project_id = ?`
args := []interface{}{projectID}
if !includeDeprecated {
query += " AND confidence != 'deprecated'"
}
query += " ORDER BY pinned DESC, updated_at DESC"
rows, err := db.Query(query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
return scanProjectFacts(rows)
}
// ListProjectFacts 分页列出项目事实。
func (db *DB) ListProjectFacts(projectID string, filter ProjectFactListFilter, limit, offset int) ([]*ProjectFact, error) {
if limit <= 0 {
limit = 100
}
query := `SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
FROM project_facts WHERE project_id = ?`
args := []interface{}{projectID}
if c := strings.TrimSpace(filter.Category); c != "" {
query += " AND category = ?"
args = append(args, c)
}
if c := strings.TrimSpace(filter.Confidence); c != "" {
query += " AND confidence = ?"
args = append(args, c)
}
if s := strings.TrimSpace(filter.Search); s != "" {
pat := "%" + s + "%"
query += " AND (fact_key LIKE ? OR summary LIKE ? OR body LIKE ?)"
args = append(args, pat, pat, pat)
}
query += " ORDER BY pinned DESC, updated_at DESC LIMIT ? OFFSET ?"
args = append(args, limit, offset)
rows, err := db.Query(query, args...)
if err != nil {
return nil, err
}
defer rows.Close()
return scanProjectFacts(rows)
}
// GetProjectFactByKey 按 key 获取事实。
func (db *DB) GetProjectFactByKey(projectID, factKey string) (*ProjectFact, error) {
row := db.QueryRow(
`SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
FROM project_facts WHERE project_id = ? AND fact_key = ?`,
projectID, factKey,
)
return scanProjectFactRow(row)
}
// GetProjectFact 按 ID 获取事实。
func (db *DB) GetProjectFact(id string) (*ProjectFact, error) {
row := db.QueryRow(
`SELECT id, project_id, fact_key, category, summary, COALESCE(body,''), confidence,
COALESCE(source_conversation_id,''), COALESCE(source_message_id,''), pinned,
COALESCE(supersedes_fact_id,''), COALESCE(related_vulnerability_id,''), created_at, updated_at
FROM project_facts WHERE id = ?`, id,
)
return scanProjectFactRow(row)
}
// UpsertProjectFact 创建或更新事实(按 project_id + fact_key)。
func (db *DB) UpsertProjectFact(f *ProjectFact) (*ProjectFact, error) {
if err := ValidateFactKey(f.FactKey); err != nil {
return nil, err
}
if strings.TrimSpace(f.Category) == "" {
f.Category = "note"
}
if strings.TrimSpace(f.Confidence) == "" {
f.Confidence = "tentative"
}
now := time.Now()
existing, err := db.GetProjectFactByKey(f.ProjectID, f.FactKey)
if err == nil && existing != nil {
f.ID = existing.ID
f.CreatedAt = existing.CreatedAt
f.UpdatedAt = now
_, err = db.Exec(
`UPDATE project_facts SET category = ?, summary = ?, body = ?, confidence = ?,
source_conversation_id = ?, source_message_id = ?, pinned = ?,
supersedes_fact_id = ?, related_vulnerability_id = ?, updated_at = ?
WHERE id = ?`,
f.Category, f.Summary, f.Body, f.Confidence,
nullIfEmpty(f.SourceConversationID), nullIfEmpty(f.SourceMessageID), boolToInt(f.Pinned),
nullIfEmpty(f.SupersedesFactID), nullIfEmpty(f.RelatedVulnerabilityID), f.UpdatedAt, f.ID,
)
if err != nil {
return nil, fmt.Errorf("更新事实失败: %w", err)
}
return f, nil
}
if f.ID == "" {
f.ID = uuid.New().String()
}
f.CreatedAt = now
f.UpdatedAt = now
_, err = db.Exec(
`INSERT INTO project_facts (
id, project_id, fact_key, category, summary, body, confidence,
source_conversation_id, source_message_id, pinned, supersedes_fact_id, related_vulnerability_id,
created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
f.ID, f.ProjectID, f.FactKey, f.Category, f.Summary, f.Body, f.Confidence,
nullIfEmpty(f.SourceConversationID), nullIfEmpty(f.SourceMessageID), boolToInt(f.Pinned),
nullIfEmpty(f.SupersedesFactID), nullIfEmpty(f.RelatedVulnerabilityID),
f.CreatedAt, f.UpdatedAt,
)
if err != nil {
return nil, fmt.Errorf("创建事实失败: %w", err)
}
return f, nil
}
// DeprecateProjectFact 将事实标记为 deprecated。
func (db *DB) DeprecateProjectFact(projectID, factKey string) error {
res, err := db.Exec(
`UPDATE project_facts SET confidence = 'deprecated', updated_at = ? WHERE project_id = ? AND fact_key = ?`,
time.Now(), projectID, factKey,
)
if err != nil {
return err
}
n, _ := res.RowsAffected()
if n == 0 {
return fmt.Errorf("事实不存在")
}
return nil
}
// DeleteProjectFact 删除事实。
func (db *DB) DeleteProjectFact(id string) error {
_, err := db.Exec(`DELETE FROM project_facts WHERE id = ?`, id)
return err
}
func scanProjectFacts(rows *sql.Rows) ([]*ProjectFact, error) {
var out []*ProjectFact
for rows.Next() {
f, err := scanProjectFactFromRows(rows)
if err != nil {
return nil, err
}
out = append(out, f)
}
return out, rows.Err()
}
func scanProjectFactRow(row *sql.Row) (*ProjectFact, error) {
var f ProjectFact
var pinned int
var createdAt, updatedAt string
err := row.Scan(
&f.ID, &f.ProjectID, &f.FactKey, &f.Category, &f.Summary, &f.Body, &f.Confidence,
&f.SourceConversationID, &f.SourceMessageID, &pinned,
&f.SupersedesFactID, &f.RelatedVulnerabilityID, &createdAt, &updatedAt,
)
if err != nil {
if err == sql.ErrNoRows {
return nil, fmt.Errorf("事实不存在")
}
return nil, err
}
f.Pinned = pinned != 0
f.CreatedAt = parseDBTime(createdAt)
f.UpdatedAt = parseDBTime(updatedAt)
return &f, nil
}
func scanProjectFactFromRows(rows *sql.Rows) (*ProjectFact, error) {
var f ProjectFact
var pinned int
var createdAt, updatedAt string
err := rows.Scan(
&f.ID, &f.ProjectID, &f.FactKey, &f.Category, &f.Summary, &f.Body, &f.Confidence,
&f.SourceConversationID, &f.SourceMessageID, &pinned,
&f.SupersedesFactID, &f.RelatedVulnerabilityID, &createdAt, &updatedAt,
)
if err != nil {
return nil, err
}
f.Pinned = pinned != 0
f.CreatedAt = parseDBTime(createdAt)
f.UpdatedAt = parseDBTime(updatedAt)
return &f, nil
}
func boolToInt(b bool) int {
if b {
return 1
}
return 0
}
func nullIfEmpty(s string) interface{} {
if strings.TrimSpace(s) == "" {
return nil
}
return s
}
func parseDBTime(s string) time.Time {
s = strings.TrimSpace(s)
if s == "" {
return time.Time{}
}
// go-sqlite3 读 DATETIME 常返回 RFC3339(含 T),写入时可能是空格分隔格式,需兼容多种形态
layouts := []string{
time.RFC3339Nano,
time.RFC3339,
"2006-01-02 15:04:05.999999999-07:00",
"2006-01-02 15:04:05-07:00",
"2006-01-02T15:04:05.999999999-07:00",
"2006-01-02T15:04:05-07:00",
"2006-01-02 15:04:05.999999999",
"2006-01-02 15:04:05",
"2006-01-02T15:04:05.999999999",
"2006-01-02T15:04:05",
}
for _, layout := range layouts {
if t, e := time.Parse(layout, s); e == nil {
return t
}
}
return time.Time{}
}
+93
View File
@@ -0,0 +1,93 @@
package database
import (
"encoding/json"
"os"
"path/filepath"
"testing"
"time"
"go.uber.org/zap"
)
func TestParseDBTime_projectFactFormats(t *testing.T) {
cases := []string{
"2026-05-26 11:13:07.442143+08:00",
"2026-05-26 11:13:07",
"2026-05-26T11:13:07.442143+08:00",
}
for _, s := range cases {
got := parseDBTime(s)
if got.IsZero() {
t.Fatalf("parseDBTime(%q) returned zero", s)
}
}
}
func TestListProjectFacts_updatedAtJSON(t *testing.T) {
root, err := os.Getwd()
if err != nil {
t.Skip(err)
}
dbPath := filepath.Join(root, "..", "..", "data", "conversations.db")
if _, err := os.Stat(dbPath); err != nil {
t.Skip("conversations.db not found")
}
db, err := NewDB(dbPath, zap.NewNop())
if err != nil {
t.Fatal(err)
}
projects, err := db.ListProjects("", 1, 0)
if err != nil || len(projects) == 0 {
t.Skip("no projects")
}
pid := projects[0].ID
list, err := db.ListProjectFacts(pid, ProjectFactListFilter{}, 5, 0)
if err != nil {
t.Fatal(err)
}
if len(list) == 0 {
t.Skip("no facts")
}
for _, f := range list {
if f.UpdatedAt.IsZero() {
t.Fatalf("fact %s UpdatedAt is zero after ListProjectFacts", f.FactKey)
}
b, err := json.Marshal(f)
if err != nil {
t.Fatal(err)
}
var m map[string]interface{}
if err := json.Unmarshal(b, &m); err != nil {
t.Fatal(err)
}
raw, ok := m["updated_at"].(string)
if !ok || raw == "" || raw[:4] == "0001" {
t.Fatalf("bad updated_at in JSON: %v", m["updated_at"])
}
}
}
func TestParseDBTime_zeroOnGarbage(t *testing.T) {
if !parseDBTime("").IsZero() {
t.Fatal("expected zero for empty")
}
}
// Ensure RFC3339 round-trip used by API is after year 2000.
func TestParseDBTime_marshalRoundTrip(t *testing.T) {
s := "2026-05-26 11:13:07.442143+08:00"
tm := parseDBTime(s)
b, err := json.Marshal(tm)
if err != nil {
t.Fatal(err)
}
var back time.Time
if err := json.Unmarshal(b, &back); err != nil {
t.Fatal(err)
}
if back.IsZero() {
t.Fatalf("unmarshal zero from %s", string(b))
}
}
+26 -9
View File
@@ -15,6 +15,7 @@ type VulnerabilityListFilter struct {
ID string ID string
Search string // 关键词模糊匹配(标题、描述、类型、目标等) Search string // 关键词模糊匹配(标题、描述、类型、目标等)
ConversationID string ConversationID string
ProjectID string
Severity string Severity string
Status string Status string
TaskID string TaskID string
@@ -38,6 +39,10 @@ func (f VulnerabilityListFilter) appendWhere(query string, args []interface{}) (
query += " AND conversation_id = ?" query += " AND conversation_id = ?"
args = append(args, f.ConversationID) args = append(args, f.ConversationID)
} }
if f.ProjectID != "" {
query += " AND project_id = ?"
args = append(args, f.ProjectID)
}
if f.TaskID != "" { if f.TaskID != "" {
query += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))" query += " AND EXISTS (SELECT 1 FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id AND (bt.id = ? OR bt.queue_id = ?))"
args = append(args, f.TaskID, f.TaskID) args = append(args, f.TaskID, f.TaskID)
@@ -85,6 +90,7 @@ func (f VulnerabilityListFilter) appendWhere(query string, args []interface{}) (
type Vulnerability struct { type Vulnerability struct {
ID string `json:"id"` ID string `json:"id"`
ConversationID string `json:"conversation_id"` ConversationID string `json:"conversation_id"`
ProjectID string `json:"project_id,omitempty"`
ConversationTag string `json:"conversation_tag,omitempty"` ConversationTag string `json:"conversation_tag,omitempty"`
TaskTag string `json:"task_tag,omitempty"` TaskTag string `json:"task_tag,omitempty"`
TaskID string `json:"task_id,omitempty"` TaskID string `json:"task_id,omitempty"`
@@ -116,17 +122,23 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
} }
vuln.UpdatedAt = now vuln.UpdatedAt = now
if strings.TrimSpace(vuln.ProjectID) == "" && vuln.ConversationID != "" {
if pid, err := db.GetConversationProjectID(vuln.ConversationID); err == nil {
vuln.ProjectID = pid
}
}
query := ` query := `
INSERT INTO vulnerabilities ( INSERT INTO vulnerabilities (
id, conversation_id, conversation_tag, task_tag, title, description, severity, status, id, conversation_id, project_id, conversation_tag, task_tag, title, description, severity, status,
vulnerability_type, target, proof, impact, recommendation, vulnerability_type, target, proof, impact, recommendation,
created_at, updated_at created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
` `
_, err := db.Exec( _, err := db.Exec(
query, query,
vuln.ID, vuln.ConversationID, vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description, vuln.ID, vuln.ConversationID, nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description,
vuln.Severity, vuln.Status, vuln.Type, vuln.Target, vuln.Severity, vuln.Status, vuln.Type, vuln.Target,
vuln.Proof, vuln.Impact, vuln.Recommendation, vuln.Proof, vuln.Impact, vuln.Recommendation,
vuln.CreatedAt, vuln.UpdatedAt, vuln.CreatedAt, vuln.UpdatedAt,
@@ -142,7 +154,7 @@ func (db *DB) CreateVulnerability(vuln *Vulnerability) (*Vulnerability, error) {
func (db *DB) GetVulnerability(id string) (*Vulnerability, error) { func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
var vuln Vulnerability var vuln Vulnerability
query := ` query := `
SELECT id, conversation_id, title, description, severity, status, SELECT id, conversation_id, COALESCE(project_id,''), title, description, severity, status,
conversation_tag, task_tag, vulnerability_type, target, proof, impact, recommendation, conversation_tag, task_tag, vulnerability_type, target, proof, impact, recommendation,
COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id, COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id, COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
@@ -152,7 +164,7 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
` `
err := db.QueryRow(query, id).Scan( err := db.QueryRow(query, id).Scan(
&vuln.ID, &vuln.ConversationID, &vuln.Title, &vuln.Description, &vuln.ID, &vuln.ConversationID, &vuln.ProjectID, &vuln.Title, &vuln.Description,
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target, &vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
&vuln.Proof, &vuln.Impact, &vuln.Recommendation, &vuln.Proof, &vuln.Impact, &vuln.Recommendation,
&vuln.TaskID, &vuln.TaskQueueID, &vuln.TaskID, &vuln.TaskQueueID,
@@ -171,7 +183,7 @@ func (db *DB) GetVulnerability(id string) (*Vulnerability, error) {
// ListVulnerabilities 列出漏洞 // ListVulnerabilities 列出漏洞
func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFilter) ([]*Vulnerability, error) { func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFilter) ([]*Vulnerability, error) {
query := ` query := `
SELECT id, conversation_id, title, description, severity, status, conversation_tag, task_tag, SELECT id, conversation_id, COALESCE(project_id,''), title, description, severity, status, conversation_tag, task_tag,
vulnerability_type, target, proof, impact, recommendation, vulnerability_type, target, proof, impact, recommendation,
COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id, COALESCE((SELECT bt.id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_id,
COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id, COALESCE((SELECT bt.queue_id FROM batch_tasks bt WHERE bt.conversation_id = vulnerabilities.conversation_id LIMIT 1), '') AS task_queue_id,
@@ -195,7 +207,7 @@ func (db *DB) ListVulnerabilities(limit, offset int, filter VulnerabilityListFil
for rows.Next() { for rows.Next() {
var vuln Vulnerability var vuln Vulnerability
err := rows.Scan( err := rows.Scan(
&vuln.ID, &vuln.ConversationID, &vuln.Title, &vuln.Description, &vuln.ID, &vuln.ConversationID, &vuln.ProjectID, &vuln.Title, &vuln.Description,
&vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target, &vuln.Severity, &vuln.Status, &vuln.ConversationTag, &vuln.TaskTag, &vuln.Type, &vuln.Target,
&vuln.Proof, &vuln.Impact, &vuln.Recommendation, &vuln.Proof, &vuln.Impact, &vuln.Recommendation,
&vuln.TaskID, &vuln.TaskQueueID, &vuln.TaskID, &vuln.TaskQueueID,
@@ -232,7 +244,7 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
query := ` query := `
UPDATE vulnerabilities UPDATE vulnerabilities
SET conversation_tag = ?, task_tag = ?, title = ?, description = ?, severity = ?, status = ?, SET project_id = ?, conversation_tag = ?, task_tag = ?, title = ?, description = ?, severity = ?, status = ?,
vulnerability_type = ?, target = ?, proof = ?, impact = ?, vulnerability_type = ?, target = ?, proof = ?, impact = ?,
recommendation = ?, updated_at = ? recommendation = ?, updated_at = ?
WHERE id = ? WHERE id = ?
@@ -240,7 +252,7 @@ func (db *DB) UpdateVulnerability(id string, vuln *Vulnerability) error {
_, err := db.Exec( _, err := db.Exec(
query, query,
vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description, vuln.Severity, vuln.Status, nullIfEmpty(vuln.ProjectID), vuln.ConversationTag, vuln.TaskTag, vuln.Title, vuln.Description, vuln.Severity, vuln.Status,
vuln.Type, vuln.Target, vuln.Proof, vuln.Impact, vuln.Type, vuln.Target, vuln.Proof, vuln.Impact,
vuln.Recommendation, vuln.UpdatedAt, id, vuln.Recommendation, vuln.UpdatedAt, id,
) )
@@ -366,10 +378,15 @@ func (db *DB) GetVulnerabilityFilterOptions() (map[string][]string, error) {
if err != nil { if err != nil {
return nil, fmt.Errorf("查询任务标签建议失败: %w", err) return nil, fmt.Errorf("查询任务标签建议失败: %w", err)
} }
projectIDs, err := collect(`SELECT DISTINCT project_id FROM vulnerabilities WHERE project_id IS NOT NULL AND project_id <> '' ORDER BY created_at DESC LIMIT 200`)
if err != nil {
return nil, fmt.Errorf("查询项目ID建议失败: %w", err)
}
return map[string][]string{ return map[string][]string{
"vulnerability_ids": vulnIDs, "vulnerability_ids": vulnIDs,
"conversation_ids": conversationIDs, "conversation_ids": conversationIDs,
"project_ids": projectIDs,
"task_ids": taskIDs, "task_ids": taskIDs,
"queue_ids": queueIDs, "queue_ids": queueIDs,
"conversation_tags": conversationTags, "conversation_tags": conversationTags,
+24 -1
View File
@@ -4,7 +4,16 @@ package builtin
// 所有代码中使用内置工具名称的地方都应该使用这些常量,而不是硬编码字符串 // 所有代码中使用内置工具名称的地方都应该使用这些常量,而不是硬编码字符串
const ( const (
// 漏洞管理工具 // 漏洞管理工具
ToolRecordVulnerability = "record_vulnerability" ToolRecordVulnerability = "record_vulnerability"
ToolListVulnerabilities = "list_vulnerabilities"
ToolGetVulnerability = "get_vulnerability"
// 项目黑板(事实)工具
ToolUpsertProjectFact = "upsert_project_fact"
ToolGetProjectFact = "get_project_fact"
ToolListProjectFacts = "list_project_facts"
ToolSearchProjectFacts = "search_project_facts"
ToolDeprecateProjectFact = "deprecate_project_fact"
// 知识库工具 // 知识库工具
ToolListKnowledgeRiskTypes = "list_knowledge_risk_types" ToolListKnowledgeRiskTypes = "list_knowledge_risk_types"
@@ -53,6 +62,13 @@ const (
func IsBuiltinTool(toolName string) bool { func IsBuiltinTool(toolName string) bool {
switch toolName { switch toolName {
case ToolRecordVulnerability, case ToolRecordVulnerability,
ToolListVulnerabilities,
ToolGetVulnerability,
ToolUpsertProjectFact,
ToolGetProjectFact,
ToolListProjectFacts,
ToolSearchProjectFacts,
ToolDeprecateProjectFact,
ToolListKnowledgeRiskTypes, ToolListKnowledgeRiskTypes,
ToolSearchKnowledgeBase, ToolSearchKnowledgeBase,
ToolWebshellExec, ToolWebshellExec,
@@ -96,6 +112,13 @@ func IsBuiltinTool(toolName string) bool {
func GetAllBuiltinTools() []string { func GetAllBuiltinTools() []string {
return []string{ return []string{
ToolRecordVulnerability, ToolRecordVulnerability,
ToolListVulnerabilities,
ToolGetVulnerability,
ToolUpsertProjectFact,
ToolGetProjectFact,
ToolListProjectFacts,
ToolSearchProjectFacts,
ToolDeprecateProjectFact,
ToolListKnowledgeRiskTypes, ToolListKnowledgeRiskTypes,
ToolSearchKnowledgeBase, ToolSearchKnowledgeBase,
ToolWebshellExec, ToolWebshellExec,