CyberStrikeAI/internal/database/vulnerability_alert_test.go

package database_test

import (
	"testing"
	"time"

	"cyberstrike-ai/internal/database"
	"cyberstrike-ai/internal/security"

	"go.uber.org/zap"
)

func TestVulnerabilityAlertSubscriptionIsOptInAndRBACScoped(t *testing.T) {
	db, err := database.NewDB(t.TempDir()+"/alerts.db", zap.NewNop())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Close() })
	if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil {
		t.Fatal(err)
	}
	owner, _ := db.CreateRBACUser("alert-owner", "Owner", "hash", true, []string{database.RBACSystemRoleOperator})
	other, _ := db.CreateRBACUser("alert-other", "Other", "hash", true, []string{database.RBACSystemRoleOperator})

	for i, pair := range []struct {
		user     *database.RBACUser
		external string
	}{{owner, "owner"}, {other, "other"}} {
		code := "code-" + string(rune('a'+i))
		if err := db.CreateRobotBindingCode(pair.user.ID, code, time.Now().Add(time.Minute)); err != nil {
			t.Fatal(err)
		}
		if _, err := db.ConsumeRobotBindingCode("wecom", pair.external, code); err != nil {
			t.Fatal(err)
		}
	}
	defaultSub, err := db.GetVulnerabilityAlertSubscription(owner.ID)
	if err != nil || defaultSub.Enabled || defaultSub.MinSeverity != "high" {
		t.Fatalf("unsafe default: %#v %v", defaultSub, err)
	}
	if _, err := db.UpsertVulnerabilityAlertSubscription(owner.ID, true, "high"); err != nil {
		t.Fatal(err)
	}
	if _, err := db.UpsertVulnerabilityAlertSubscription(other.ID, true, "low"); err != nil {
		t.Fatal(err)
	}

	vuln, err := db.CreateVulnerability(&database.Vulnerability{Title: "owned", Severity: "high"})
	if err != nil {
		t.Fatal(err)
	}
	if err := db.SetResourceOwner("vulnerability", vuln.ID, owner.ID); err != nil {
		t.Fatal(err)
	}
	recipients, err := db.ListVulnerabilityAlertRecipients(vuln)
	if err != nil {
		t.Fatal(err)
	}
	if len(recipients) != 1 || recipients[0].UserID != owner.ID || recipients[0].ExternalUserID != "owner" {
		t.Fatalf("alert escaped RBAC boundary: %#v", recipients)
	}
	if err := db.EnqueueVulnerabilityAlertDeliveries(vuln.ID, recipients); err != nil {
		t.Fatal(err)
	}
	if err := db.EnqueueVulnerabilityAlertDeliveries(vuln.ID, recipients); err != nil {
		t.Fatal(err)
	}
	deliveries, err := db.ListDueVulnerabilityAlertDeliveries(10)
	if err != nil || len(deliveries) != 1 {
		t.Fatalf("outbox is not durable/deduplicated: %#v %v", deliveries, err)
	}
	if err := db.MarkVulnerabilityAlertDeliverySent(deliveries[0].ID); err != nil {
		t.Fatal(err)
	}
	deliveries, _ = db.ListDueVulnerabilityAlertDeliveries(10)
	if len(deliveries) != 0 {
		t.Fatalf("sent delivery remained due: %#v", deliveries)
	}

	vuln.Severity = "medium"
	recipients, err = db.ListVulnerabilityAlertRecipients(vuln)
	if err != nil {
		t.Fatal(err)
	}
	if len(recipients) != 0 {
		t.Fatalf("minimum severity was ignored: %#v", recipients)
	}
}