CyberStrikeAI/internal/database/vulnerability_alert.go

package database

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

// VulnerabilityAlertSubscription is the single source of truth shared by Web
// settings and robot commands. Alerts are opt-in and user scoped.
type VulnerabilityAlertSubscription struct {
	UserID      string    `json:"user_id"`
	Enabled     bool      `json:"enabled"`
	MinSeverity string    `json:"min_severity"`
	CreatedAt   time.Time `json:"created_at"`
	UpdatedAt   time.Time `json:"updated_at"`
}

type VulnerabilityAlertRecipient struct {
	UserID         string
	Platform       string
	ExternalUserID string
}

type VulnerabilityAlertDelivery struct {
	ID             int64
	Vulnerability  *Vulnerability
	UserID         string
	Platform       string
	ExternalUserID string
	Attempts       int
}

var vulnerabilitySeverityRank = map[string]int{
	"info": 0, "low": 1, "medium": 2, "high": 3, "critical": 4,
}

func NormalizeVulnerabilityAlertSeverity(value string) (string, error) {
	value = strings.ToLower(strings.TrimSpace(value))
	if value == "" {
		value = "high"
	}
	if _, ok := vulnerabilitySeverityRank[value]; !ok {
		return "", fmt.Errorf("invalid minimum severity %q", value)
	}
	return value, nil
}

func (db *DB) GetVulnerabilityAlertSubscription(userID string) (*VulnerabilityAlertSubscription, error) {
	userID = strings.TrimSpace(userID)
	var sub VulnerabilityAlertSubscription
	var enabled int
	var createdAt, updatedAt string
	err := db.QueryRow(`SELECT user_id, enabled, min_severity, created_at, updated_at
		FROM vulnerability_alert_subscriptions WHERE user_id = ?`, userID).
		Scan(&sub.UserID, &enabled, &sub.MinSeverity, &createdAt, &updatedAt)
	if err == sql.ErrNoRows {
		now := time.Now()
		return &VulnerabilityAlertSubscription{UserID: userID, MinSeverity: "high", CreatedAt: now, UpdatedAt: now}, nil
	}
	if err != nil {
		return nil, err
	}
	sub.Enabled = enabled != 0
	sub.CreatedAt = parseDBTime(createdAt)
	sub.UpdatedAt = parseDBTime(updatedAt)
	return &sub, nil
}

func (db *DB) UpsertVulnerabilityAlertSubscription(userID string, enabled bool, minSeverity string) (*VulnerabilityAlertSubscription, error) {
	userID = strings.TrimSpace(userID)
	if userID == "" {
		return nil, fmt.Errorf("user id is required")
	}
	severity, err := NormalizeVulnerabilityAlertSeverity(minSeverity)
	if err != nil {
		return nil, err
	}
	now := time.Now()
	_, err = db.Exec(`INSERT INTO vulnerability_alert_subscriptions
		(user_id, enabled, min_severity, created_at, updated_at) VALUES (?, ?, ?, ?, ?)
		ON CONFLICT(user_id) DO UPDATE SET enabled = excluded.enabled,
		min_severity = excluded.min_severity, updated_at = excluded.updated_at`,
		userID, boolToInt(enabled), severity, now, now)
	if err != nil {
		return nil, err
	}
	return db.GetVulnerabilityAlertSubscription(userID)
}

// ListVulnerabilityAlertRecipients applies the same RBAC ownership/assignment
// boundaries as the vulnerability list, then expands only enabled robot bindings.
func (db *DB) ListVulnerabilityAlertRecipients(vuln *Vulnerability) ([]VulnerabilityAlertRecipient, error) {
	if vuln == nil {
		return nil, nil
	}
	rank, ok := vulnerabilitySeverityRank[strings.ToLower(strings.TrimSpace(vuln.Severity))]
	if !ok {
		return nil, nil
	}
	rows, err := db.Query(`
		SELECT DISTINCT s.user_id, b.platform, b.external_user_id, s.min_severity
		FROM vulnerability_alert_subscriptions s
		JOIN rbac_users u ON u.id = s.user_id AND u.enabled = 1
		JOIN robot_user_bindings b ON b.rbac_user_id = s.user_id AND b.enabled = 1
		WHERE s.enabled = 1 AND (
			EXISTS (SELECT 1 FROM vulnerabilities v WHERE v.id = ? AND v.owner_user_id = s.user_id)
			OR EXISTS (SELECT 1 FROM rbac_resource_assignments ra WHERE ra.user_id = s.user_id AND ra.resource_type = 'vulnerability' AND ra.resource_id = ?)
			OR (? <> '' AND (EXISTS (SELECT 1 FROM projects p WHERE p.id = ? AND p.owner_user_id = s.user_id)
				OR EXISTS (SELECT 1 FROM rbac_resource_assignments pra WHERE pra.user_id = s.user_id AND pra.resource_type = 'project' AND pra.resource_id = ?)))
			OR (? <> '' AND (EXISTS (SELECT 1 FROM conversations c WHERE c.id = ? AND c.owner_user_id = s.user_id)
				OR EXISTS (SELECT 1 FROM rbac_resource_assignments cra WHERE cra.user_id = s.user_id AND cra.resource_type = 'conversation' AND cra.resource_id = ?)))
		)`, vuln.ID, vuln.ID, vuln.ProjectID, vuln.ProjectID, vuln.ProjectID,
		vuln.ConversationID, vuln.ConversationID, vuln.ConversationID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := make([]VulnerabilityAlertRecipient, 0)
	for rows.Next() {
		var recipient VulnerabilityAlertRecipient
		var minimum string
		if err := rows.Scan(&recipient.UserID, &recipient.Platform, &recipient.ExternalUserID, &minimum); err != nil {
			return nil, err
		}
		if rank >= vulnerabilitySeverityRank[minimum] {
			out = append(out, recipient)
		}
	}
	return out, rows.Err()
}

func (db *DB) SetVulnerabilityCreatedHook(hook func(*Vulnerability)) {
	db.vulnerabilityCreatedHook = hook
}

// NotifyVulnerabilityCreated must be called after resource ownership has been
// committed. Delivery runs asynchronously and never delays the write path.
func (db *DB) NotifyVulnerabilityCreated(vulnerability *Vulnerability) {
	if db == nil || vulnerability == nil || db.vulnerabilityCreatedHook == nil {
		return
	}
	created := *vulnerability
	go db.vulnerabilityCreatedHook(&created)
}

func (db *DB) EnqueueVulnerabilityAlertDeliveries(vulnerabilityID string, recipients []VulnerabilityAlertRecipient) error {
	now := time.Now()
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()
	for _, r := range recipients {
		if _, err := tx.Exec(`INSERT INTO vulnerability_alert_deliveries
			(vulnerability_id, user_id, platform, external_user_id, status, attempts, next_attempt_at, created_at, updated_at)
			VALUES (?, ?, ?, ?, 'pending', 0, ?, ?, ?)
			ON CONFLICT(vulnerability_id, platform, external_user_id) DO NOTHING`,
			vulnerabilityID, r.UserID, r.Platform, r.ExternalUserID, now, now, now); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func (db *DB) ListDueVulnerabilityAlertDeliveries(limit int) ([]VulnerabilityAlertDelivery, error) {
	if limit <= 0 || limit > 100 {
		limit = 50
	}
	rows, err := db.Query(`SELECT d.id, d.user_id, d.platform, d.external_user_id, d.attempts,
		v.id, COALESCE(v.conversation_id,''), COALESCE(v.project_id,''), v.title, COALESCE(v.description,''),
		v.severity, v.status, COALESCE(v.vulnerability_type,''), COALESCE(v.target,''),
		COALESCE(v.impact,''), COALESCE(v.recommendation,''), v.created_at, v.updated_at
		FROM vulnerability_alert_deliveries d JOIN vulnerabilities v ON v.id = d.vulnerability_id
		WHERE d.status IN ('pending','retry') AND d.next_attempt_at <= ?
		ORDER BY d.next_attempt_at, d.id LIMIT ?`, time.Now(), limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []VulnerabilityAlertDelivery
	for rows.Next() {
		var d VulnerabilityAlertDelivery
		v := &Vulnerability{}
		if err := rows.Scan(&d.ID, &d.UserID, &d.Platform, &d.ExternalUserID, &d.Attempts,
			&v.ID, &v.ConversationID, &v.ProjectID, &v.Title, &v.Description, &v.Severity, &v.Status,
			&v.Type, &v.Target, &v.Impact, &v.Recommendation, &v.CreatedAt, &v.UpdatedAt); err != nil {
			return nil, err
		}
		d.Vulnerability = v
		out = append(out, d)
	}
	return out, rows.Err()
}

func (db *DB) MarkVulnerabilityAlertDeliverySent(id int64) error {
	_, err := db.Exec(`UPDATE vulnerability_alert_deliveries SET status='sent', attempts=attempts+1, last_error='', updated_at=? WHERE id=?`, time.Now(), id)
	return err
}

func (db *DB) MarkVulnerabilityAlertDeliveryFailed(id int64, attempts int, sendErr error) error {
	status := "retry"
	if attempts >= 5 {
		status = "failed"
	}
	delay := time.Minute * time.Duration(1<<min(attempts, 6))
	message := ""
	if sendErr != nil {
		message = sendErr.Error()
	}
	_, err := db.Exec(`UPDATE vulnerability_alert_deliveries SET status=?, attempts=?, next_attempt_at=?, last_error=?, updated_at=? WHERE id=?`,
		status, attempts, time.Now().Add(delay), message, time.Now(), id)
	return err
}