CyberStrikeAI/internal/handler/vulnerability_alert_test.go

package handler

import (
	"path/filepath"
	"strings"
	"testing"
	"time"

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

	"go.uber.org/zap"
)

func TestRobotVulnerabilityAlertCommandSharesSubscription(t *testing.T) {
	db, err := database.NewDB(filepath.Join(t.TempDir(), "alert-command.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)
	}
	user, err := db.CreateRBACUser("alert-user", "Alert User", "hash", true, []string{database.RBACSystemRoleOperator})
	if err != nil {
		t.Fatal(err)
	}
	if err := db.CreateRobotBindingCode(user.ID, "alert-code", time.Now().Add(time.Minute)); err != nil {
		t.Fatal(err)
	}
	if _, err := db.ConsumeRobotBindingCode("wecom", "external-user", "alert-code"); err != nil {
		t.Fatal(err)
	}
	h := NewRobotHandler(&config.Config{}, db, nil, zap.NewNop())

	if got := h.HandleMessage("wecom", "external-user", "漏洞提醒 高危以上"); !strings.Contains(got, "已开启") {
		t.Fatalf("enable reply: %s", got)
	}
	sub, err := db.GetVulnerabilityAlertSubscription(user.ID)
	if err != nil || !sub.Enabled || sub.MinSeverity != "high" {
		t.Fatalf("web subscription not updated: %#v %v", sub, err)
	}
	if got := h.HandleMessage("wecom", "external-user", "vuln alerts off"); !strings.Contains(got, "已关闭") {
		t.Fatalf("disable reply: %s", got)
	}
	sub, _ = db.GetVulnerabilityAlertSubscription(user.ID)
	if sub.Enabled {
		t.Fatalf("subscription remained enabled: %#v", sub)
	}
}