CyberStrikeAI/internal/handler/rest_semantics_test.go

package handler

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"
	"time"

	"cyberstrike-ai/internal/c2"
	"cyberstrike-ai/internal/database"
	"github.com/gin-gonic/gin"
	"go.uber.org/zap"
)

func newRESTSemanticsDB(t *testing.T) *database.DB {
	t.Helper()
	db, err := database.NewDB(filepath.Join(t.TempDir(), "rest.db"), zap.NewNop())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Close() })
	return db
}

func restSemanticsRequest(r http.Handler, method, path string) *httptest.ResponseRecorder {
	w := httptest.NewRecorder()
	r.ServeHTTP(w, httptest.NewRequest(method, path, nil))
	return w
}

func TestRESTDeleteSemantics(t *testing.T) {
	for _, module := range []string{"profile", "webshell"} {
		t.Run(module, func(t *testing.T) {
			db := newRESTSemanticsDB(t)
			r := gin.New()
			table, missing, failure, success := "c2_profiles", "profile not found", "failed to delete profile", `{"deleted":true}`
			if module == "profile" {
				h := NewC2Handler(c2.NewManager(db, zap.NewNop(), t.TempDir()), zap.NewNop())
				r.DELETE("/objects/:id", h.DeleteProfile)
				if err := db.CreateC2Profile(&database.C2Profile{ID: "existing", Name: "test", CreatedAt: time.Now()}); err != nil {
					t.Fatal(err)
				}
			} else {
				table, missing, failure, success = "webshell_connections", "connection not found", "failed to delete connection", `{"ok":true}`
				h := NewWebShellHandler(zap.NewNop(), db)
				r.DELETE("/objects/:id", h.DeleteConnection)
				if err := db.CreateWebshellConnection(&database.WebShellConnection{ID: "existing", URL: "https://example.invalid", CreatedAt: time.Now()}); err != nil {
					t.Fatal(err)
				}
			}
			for _, tc := range []struct {
				id   string
				code int
				body string
			}{
				{"missing", http.StatusNotFound, `{"error":"` + missing + `"}`},
				{"existing", http.StatusOK, success},
				{"existing", http.StatusNotFound, `{"error":"` + missing + `"}`},
			} {
				w := restSemanticsRequest(r, http.MethodDelete, "/objects/"+tc.id)
				if w.Code != tc.code || w.Body.String() != tc.body {
					t.Fatalf("delete %s: %d %s, want %d %s", tc.id, w.Code, w.Body.String(), tc.code, tc.body)
				}
			}
			// Break only the temporary database: genuine SQL failures must remain 500,
			// with no table name or driver error in the HTTP response.
			if _, err := db.Exec("DROP TABLE " + table); err != nil {
				t.Fatal(err)
			}
			w := restSemanticsRequest(r, http.MethodDelete, "/objects/missing")
			if w.Code != http.StatusInternalServerError || w.Body.String() != `{"error":"`+failure+`"}` {
				t.Fatalf("SQL failure: %d %s", w.Code, w.Body.String())
			}
		})
	}
}

func assertRESTList(t *testing.T, r http.Handler, path, field string, count int) {
	t.Helper()
	w := restSemanticsRequest(r, http.MethodGet, path)
	if w.Code != http.StatusOK {
		t.Fatalf("%s: %d %s", path, w.Code, w.Body.String())
	}
	var body map[string]json.RawMessage
	if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
		t.Fatal(err)
	}
	var items []json.RawMessage
	if err := json.Unmarshal(body[field], &items); err != nil || items == nil || len(items) != count {
		t.Fatalf("%s: %s=%s, want array of length %d (err=%v)", path, field, body[field], count, err)
	}
	var total int
	if err := json.Unmarshal(body["total"], &total); err != nil || total != count {
		t.Fatalf("%s: total=%s, want %d", path, body["total"], count)
	}
}

func TestRESTVulnerabilityListSemantics(t *testing.T) {
	db := newRESTSemanticsDB(t)
	h := NewVulnerabilityHandler(db, zap.NewNop())
	r := gin.New()
	r.GET("/api/vulnerabilities", h.ListVulnerabilities)
	assertRESTList(t, r, "/api/vulnerabilities", "vulnerabilities", 0)
	if _, err := db.CreateVulnerability(&database.Vulnerability{Title: "test finding", Status: "open", Severity: "low"}); err != nil {
		t.Fatal(err)
	}
	assertRESTList(t, r, "/api/vulnerabilities?status=__none__", "vulnerabilities", 0)
	assertRESTList(t, r, "/api/vulnerabilities?status=open", "vulnerabilities", 1)
	if _, err := db.Exec("DROP TABLE vulnerabilities"); err != nil {
		t.Fatal(err)
	}
	w := restSemanticsRequest(r, http.MethodGet, "/api/vulnerabilities")
	if w.Code != http.StatusInternalServerError {
		t.Fatalf("SQL failure became empty success: %d %s", w.Code, w.Body.String())
	}
}

func TestRESTBatchListSemantics(t *testing.T) {
	for _, withDB := range []bool{true, false} {
		name := "memory"
		if withDB {
			name = "database"
		}
		t.Run(name, func(t *testing.T) {
			m := NewBatchTaskManager(zap.NewNop())
			var db *database.DB
			if withDB {
				db = newRESTSemanticsDB(t)
				m.SetDB(db)
			}
			h := &AgentHandler{batchTaskManager: m, logger: zap.NewNop()}
			r := gin.New()
			r.GET("/api/batch-tasks", h.ListBatchQueues)
			assertRESTList(t, r, "/api/batch-tasks", "queues", 0)
			if _, err := m.CreateBatchQueue("test queue", "", "eino_single", "manual", "", "", nil, 1, []string{"test"}); err != nil {
				t.Fatal(err)
			}
			assertRESTList(t, r, "/api/batch-tasks?status=__none__", "queues", 0)
			assertRESTList(t, r, "/api/batch-tasks", "queues", 1)
			if withDB {
				if _, err := db.Exec("DROP TABLE batch_task_queues"); err != nil {
					t.Fatal(err)
				}
				w := restSemanticsRequest(r, http.MethodGet, "/api/batch-tasks")
				if w.Code != http.StatusInternalServerError {
					t.Fatalf("SQL failure became empty success: %d %s", w.Code, w.Body.String())
				}
			}
		})
	}
}