package handler import ( "bytes" "encoding/json" "net/http" "net/http/httptest" "path/filepath" "testing" "time" "cyberstrike-ai/internal/database" "cyberstrike-ai/internal/security" "github.com/gin-gonic/gin" "go.uber.org/zap" ) func TestWebshellExecRequiresConnectionAccessWhenConnectionIDProvided(t *testing.T) { gin.SetMode(gin.TestMode) db, user, allowed, hidden := setupWebshellRBACTest(t) handler := NewWebShellHandler(zap.NewNop(), db) w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{ "url": hidden.URL, "connection_id": hidden.ID, "command": "id", }, handler.Exec) if w.Code != http.StatusForbidden { t.Fatalf("hidden connection status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) } w = performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{ "url": hidden.URL, "connection_id": allowed.ID, "command": "id", }, handler.Exec) if w.Code != http.StatusForbidden { t.Fatalf("mismatched URL status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) } } func TestWebshellExecRejectsAdHocURLWithoutConnectionID(t *testing.T) { gin.SetMode(gin.TestMode) _, user, _, _ := setupWebshellRBACTest(t) handler := NewWebShellHandler(zap.NewNop(), nil) w := performWebshellJSON(user, http.MethodPost, "/api/webshell/exec", map[string]interface{}{ "url": "http://127.0.0.1/admin", "command": "id", }, handler.Exec) if w.Code != http.StatusForbidden { t.Fatalf("ad-hoc URL status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) } } func TestWebshellFileOpRequiresConnectionAccessWhenConnectionIDProvided(t *testing.T) { gin.SetMode(gin.TestMode) db, user, _, hidden := setupWebshellRBACTest(t) handler := NewWebShellHandler(zap.NewNop(), db) w := performWebshellJSON(user, http.MethodPost, "/api/webshell/file", map[string]interface{}{ "url": hidden.URL, "connection_id": hidden.ID, "action": "list", "path": ".", }, handler.FileOp) if w.Code != http.StatusForbidden { t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusForbidden, w.Body.String()) } } func setupWebshellRBACTest(t *testing.T) (*database.DB, *database.RBACUser, *database.WebShellConnection, *database.WebShellConnection) { t.Helper() db, err := database.NewDB(filepath.Join(t.TempDir(), "webshell-rbac.db"), zap.NewNop()) if err != nil { t.Fatalf("NewDB: %v", err) } t.Cleanup(func() { _ = db.Close() }) user, err := db.CreateRBACUser("operator1", "Operator One", "hash", true, nil) if err != nil { t.Fatalf("CreateRBACUser: %v", err) } allowed := &database.WebShellConnection{ ID: "ws_allowed", URL: "http://127.0.0.1/allowed.php", Type: "php", Method: "post", CmdParam: "cmd", CreatedAt: time.Now(), } hidden := &database.WebShellConnection{ ID: "ws_hidden", URL: "http://127.0.0.1/hidden.php", Type: "php", Method: "post", CmdParam: "cmd", CreatedAt: time.Now(), } if err := db.CreateWebshellConnection(allowed); err != nil { t.Fatalf("CreateWebshellConnection allowed: %v", err) } if err := db.CreateWebshellConnection(hidden); err != nil { t.Fatalf("CreateWebshellConnection hidden: %v", err) } if err := db.AssignResourceToUser(user.ID, "webshell", allowed.ID); err != nil { t.Fatalf("AssignResourceToUser: %v", err) } return db, user, allowed, hidden } func performWebshellJSON(user *database.RBACUser, method, path string, body map[string]interface{}, handler gin.HandlerFunc) *httptest.ResponseRecorder { payload, _ := json.Marshal(body) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Request = httptest.NewRequest(method, path, bytes.NewReader(payload)) c.Request.Header.Set("Content-Type", "application/json") c.Set(security.ContextSessionKey, security.Session{ UserID: user.ID, Username: user.Username, Permissions: map[string]bool{"webshell:write": true}, Scope: database.RBACScopeAssigned, }) handler(c) return w }