CyberStrikeAI/internal/database/vulnerability_retest.go

package database

import (
	"database/sql"
	"fmt"
	"strings"
	"time"

	"github.com/google/uuid"
)

// VulnerabilityRetest identifies the independent snapshot used for a retest.
type VulnerabilityRetest struct {
	ConversationID       string `json:"conversation_id"`
	SourceConversationID string `json:"source_conversation_id"`
	VulnerabilityID      string `json:"vulnerability_id"`
	SourceProjectID      string `json:"source_project_id"`
	PendingPrompt        string `json:"pending_prompt,omitempty"`
}

func (db *DB) GetVulnerabilityRetest(conversationID string) (*VulnerabilityRetest, error) {
	r := &VulnerabilityRetest{ConversationID: conversationID}
	err := db.QueryRow(`SELECT source_conversation_id, vulnerability_id, source_project_id, pending_prompt FROM vulnerability_retests WHERE conversation_id = ?`, conversationID).Scan(&r.SourceConversationID, &r.VulnerabilityID, &r.SourceProjectID, &r.PendingPrompt)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	return r, err
}

// ResolveConversationPlaceholders replaces identifiers throughout the input,
// including follow-up turns. Unknown placeholders remain untouched.
func (db *DB) ResolveConversationPlaceholders(conversationID, input string) (string, error) {
	r, err := db.GetVulnerabilityRetest(conversationID)
	if err != nil {
		return "", err
	}
	replacements := []string{"{{conversation_id}}", conversationID}
	if r != nil {
		replacements = append(replacements, "{{source_conversation_id}}", r.SourceConversationID, "{{vulnerability_id}}", r.VulnerabilityID)
	}
	return strings.NewReplacer(replacements...).Replace(input), nil
}

// copiedTime re-binds a scanned timestamp as time.Time so copied rows use the same created_at format AddMessage writes, which the raw string ORDER BY depends on.
func copiedTime(value sql.NullString) interface{} {
	if !value.Valid || strings.TrimSpace(value.String) == "" {
		return nil
	}
	for _, layout := range []string{time.RFC3339Nano, time.RFC3339, "2006-01-02 15:04:05.999999999-07:00", "2006-01-02 15:04:05"} {
		if parsed, err := time.Parse(layout, value.String); err == nil {
			return parsed
		}
	}
	return value.String
}

// CreateVulnerabilityRetest snapshots history and model trace in one transaction.
// Each call creates a sibling; nothing is written to the source conversation.
func (db *DB) CreateVulnerabilityRetest(v *Vulnerability, meta ConversationCreateMeta) (*Conversation, error) {
	tx, err := db.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()
	id, now := uuid.NewString(), time.Now()
	title := "复测：" + v.Title
	result, err := tx.Exec(`INSERT INTO conversations (id, title, project_id, role_name, agent_mode, last_react_input, last_react_output, created_at, updated_at)
 SELECT ?, ?, project_id, role_name, agent_mode, last_react_input, last_react_output, ?, ? FROM conversations WHERE id = ?`, id, title, now, now, v.ConversationID)
	if err != nil {
		return nil, err
	}
	n, err := result.RowsAffected()
	if err != nil || n != 1 {
		return nil, fmt.Errorf("原始对话不存在")
	}
	conv := &Conversation{ID: id, Title: title, CreatedAt: now, UpdatedAt: now}
	if err = tx.QueryRow(`SELECT COALESCE(project_id,''), role_name, agent_mode FROM conversations WHERE id=?`, id).Scan(&conv.ProjectID, &conv.RoleName, &conv.AgentMode); err != nil {
		return nil, err
	}
	// Copy visible messages using new IDs, preserving chronological order. Execution
	// records belong to the source and must not be shared with a mutable branch.
	rows, err := tx.Query(`SELECT role, content, reasoning_content, created_at, updated_at FROM messages WHERE conversation_id=? ORDER BY created_at, rowid`, v.ConversationID)
	if err != nil {
		return nil, err
	}
	type snapshotMessage struct {
		role, content               string
		reasoning, created, updated sql.NullString
	}
	var messages []snapshotMessage
	for rows.Next() {
		var m snapshotMessage
		if err = rows.Scan(&m.role, &m.content, &m.reasoning, &m.created, &m.updated); err != nil {
			rows.Close()
			return nil, err
		}
		messages = append(messages, m)
	}
	err = rows.Err()
	rows.Close()
	if err != nil {
		return nil, err
	}
	for _, m := range messages {
		if _, err = tx.Exec(`INSERT INTO messages (id, conversation_id, role, content, reasoning_content, created_at, updated_at) VALUES (?,?,?,?,?,?,?)`, uuid.NewString(), id, m.role, m.content, m.reasoning, copiedTime(m.created), copiedTime(m.updated)); err != nil {
			return nil, err
		}
	}
	if _, err = tx.Exec(`INSERT INTO vulnerability_retests (conversation_id, source_conversation_id, vulnerability_id, source_project_id) VALUES (?,?,?,?)`, id, v.ConversationID, v.ID, conv.ProjectID); err != nil {
		return nil, err
	}
	// Preserve the source's approval policy without copying live interrupts.
	// HITL schema is installed by the agent handler rather than by NewDB.
	var hasHITL int
	if err = tx.QueryRow(`SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='hitl_conversation_configs'`).Scan(&hasHITL); err != nil {
		return nil, err
	}
	if hasHITL != 0 {
		if _, err = tx.Exec(`INSERT INTO hitl_conversation_configs (conversation_id, enabled, mode, reviewer, sensitive_tools, timeout_seconds, updated_at)
            SELECT ?, enabled, mode, reviewer, sensitive_tools, timeout_seconds, ? FROM hitl_conversation_configs WHERE conversation_id=?`, id, now, v.ConversationID); err != nil {
			return nil, err
		}
	}
	if err = tx.Commit(); err != nil {
		return nil, err
	}
	meta.ProjectID, meta.RoleName, meta.AgentMode = conv.ProjectID, conv.RoleName, conv.AgentMode
	notifyConversationCreated(conv, meta)
	return conv, nil
}

// UpdateVulnerabilityRetestResult appends evidence to retest_log; retest_notes stays operator-owned.
func (db *DB) UpdateVulnerabilityRetestResult(id, status, notes, conversationID string) error {
	entry := fmt.Sprintf("\n\n[复测 %s | 对话 %s | 状态 %s]\n%s", time.Now().Format(time.RFC3339), conversationID, status, notes)
	result, err := db.Exec(`UPDATE vulnerabilities SET status=?, retest_log=COALESCE(retest_log,'') || ?, updated_at=? WHERE id=?`, status, entry, time.Now(), id)
	if err != nil {
		return err
	}
	n, err := result.RowsAffected()
	if err != nil {
		return err
	}
	if n != 1 {
		return fmt.Errorf("漏洞不存在")
	}
	if v, err := db.GetVulnerability(id); err == nil {
		db.refreshAssetRiskCacheForConversationsBestEffort(v.ConversationID)
	}
	return nil
}

// A pending prompt survives a lost response, navigation or browser refresh.
func (db *DB) SetVulnerabilityRetestPrompt(conversationID, prompt string) error {
	_, err := db.Exec(`UPDATE vulnerability_retests SET pending_prompt=? WHERE conversation_id=?`, prompt, conversationID)
	return err
}