mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-10-18 21:47:17 +02:00
103 lines
3.6 KiB
Go
103 lines
3.6 KiB
Go
package handler
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
|
|
"cyberstrike-ai/internal/database"
|
|
"cyberstrike-ai/internal/security"
|
|
"github.com/gin-gonic/gin"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
func TestManualCreateVulnerability(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
db, err := database.NewDB(filepath.Join(t.TempDir(), "manual.db"), zap.NewNop())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer db.Close()
|
|
project, err := db.CreateProject(&database.Project{Name: "manual project"})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
conversation, err := db.CreateConversation("source", database.ConversationCreateMeta{ProjectID: project.ID})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
user, err := db.CreateRBACUser("manual-reporter", "Reporter", "hash", true, nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
h := NewVulnerabilityHandler(db, zap.NewNop())
|
|
for _, tc := range []struct {
|
|
name, conversationID, projectID, scope string
|
|
code int
|
|
wantConversation, wantProject, errorText string
|
|
}{
|
|
{name: "standalone", code: 200},
|
|
{name: "whitespace standalone", conversationID: " ", projectID: " ", code: 200},
|
|
{name: "project only", projectID: project.ID, code: 200, wantProject: project.ID},
|
|
{name: "existing conversation inherits project", conversationID: " " + conversation.ID + " ", code: 200, wantConversation: conversation.ID, wantProject: project.ID},
|
|
{name: "missing conversation", conversationID: "missing", code: 400, errorText: "会话不存在"},
|
|
{name: "missing project", projectID: "missing", code: 400, errorText: "项目不存在"},
|
|
{name: "assigned standalone", scope: database.RBACScopeAssigned, code: 200},
|
|
{name: "inaccessible conversation", conversationID: conversation.ID, scope: database.RBACScopeAssigned, code: 403},
|
|
{name: "inaccessible project", projectID: project.ID, scope: database.RBACScopeAssigned, code: 403},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
payload := map[string]string{"title": "manual finding", "severity": "high", "conversation_id": tc.conversationID, "project_id": tc.projectID}
|
|
if tc.name == "standalone" {
|
|
delete(payload, "conversation_id")
|
|
}
|
|
body, _ := json.Marshal(payload)
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/api/vulnerabilities", bytes.NewReader(body))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
scope := tc.scope
|
|
if scope == "" {
|
|
scope = database.RBACScopeAll
|
|
}
|
|
c.Set(security.ContextSessionKey, security.Session{UserID: user.ID, Scope: scope})
|
|
var before int
|
|
if err := db.QueryRow("SELECT COUNT(*) FROM vulnerabilities").Scan(&before); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
h.CreateVulnerability(c)
|
|
if w.Code != tc.code || !strings.Contains(w.Body.String(), tc.errorText) {
|
|
t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if tc.code != http.StatusOK {
|
|
var after int
|
|
if err := db.QueryRow("SELECT COUNT(*) FROM vulnerabilities").Scan(&after); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if before != after {
|
|
t.Fatal("invalid finding persisted")
|
|
}
|
|
return
|
|
}
|
|
var created database.Vulnerability
|
|
if err := json.Unmarshal(w.Body.Bytes(), &created); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
stored, err := db.GetVulnerability(created.ID)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if stored.ConversationID != tc.wantConversation || stored.ProjectID != tc.wantProject {
|
|
t.Fatalf("unexpected links: %+v", stored)
|
|
}
|
|
if !db.UserCanAccessResource(user.ID, database.RBACScopeAssigned, "vulnerability", created.ID) {
|
|
t.Fatal("creator cannot access standalone finding")
|
|
}
|
|
})
|
|
}
|
|
}
|