merge: phase 18 API+SSE
This commit is contained in:
@@ -56,7 +56,6 @@ CREATE TABLE IF NOT EXISTS custom_dorks (
|
||||
CREATE INDEX IF NOT EXISTS idx_custom_dorks_source ON custom_dorks(source);
|
||||
CREATE INDEX IF NOT EXISTS idx_custom_dorks_category ON custom_dorks(category);
|
||||
|
||||
<<<<<<< HEAD
|
||||
-- Phase 17: Telegram bot subscribers for auto-notifications.
|
||||
CREATE TABLE IF NOT EXISTS subscribers (
|
||||
chat_id INTEGER PRIMARY KEY,
|
||||
@@ -76,19 +75,5 @@ CREATE TABLE IF NOT EXISTS scheduled_jobs (
|
||||
next_run DATETIME,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
=======
|
||||
-- Phase 17: scheduled scan jobs for cron-based recurring scans.
|
||||
CREATE TABLE IF NOT EXISTS scheduled_jobs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
name TEXT NOT NULL,
|
||||
cron_expr TEXT NOT NULL,
|
||||
scan_path TEXT NOT NULL,
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
notify INTEGER NOT NULL DEFAULT 1,
|
||||
last_run_at DATETIME,
|
||||
next_run_at DATETIME,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_scheduled_jobs_enabled ON scheduled_jobs(enabled);
|
||||
>>>>>>> worktree-agent-a39573e4
|
||||
|
||||
481
pkg/web/api.go
Normal file
481
pkg/web/api.go
Normal file
@@ -0,0 +1,481 @@
|
||||
package web
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/salvacybersec/keyhunter/pkg/dorks"
|
||||
"github.com/salvacybersec/keyhunter/pkg/providers"
|
||||
"github.com/salvacybersec/keyhunter/pkg/recon"
|
||||
"github.com/salvacybersec/keyhunter/pkg/storage"
|
||||
"github.com/spf13/viper"
|
||||
)
|
||||
|
||||
// apiKey is the JSON-friendly representation of a finding with masked key value.
|
||||
type apiKey struct {
|
||||
ID int64 `json:"id"`
|
||||
ScanID int64 `json:"scanId,omitempty"`
|
||||
ProviderName string `json:"providerName"`
|
||||
KeyValue string `json:"keyValue"`
|
||||
KeyMasked string `json:"keyMasked"`
|
||||
Confidence string `json:"confidence"`
|
||||
SourcePath string `json:"sourcePath"`
|
||||
SourceType string `json:"sourceType"`
|
||||
LineNumber int `json:"lineNumber"`
|
||||
Verified bool `json:"verified"`
|
||||
VerifyStatus string `json:"verifyStatus,omitempty"`
|
||||
VerifyHTTPCode int `json:"verifyHttpCode,omitempty"`
|
||||
VerifyMetadata map[string]string `json:"verifyMetadata,omitempty"`
|
||||
CreatedAt string `json:"createdAt"`
|
||||
}
|
||||
|
||||
// mountAPI registers all /api/v1/* routes on the given router.
|
||||
func (s *Server) mountAPI(r chi.Router) {
|
||||
r.Route("/api/v1", func(r chi.Router) {
|
||||
// Stats
|
||||
r.Get("/stats", s.handleAPIStats)
|
||||
|
||||
// Keys
|
||||
r.Get("/keys", s.handleAPIListKeys)
|
||||
r.Get("/keys/{id}", s.handleAPIGetKey)
|
||||
r.Delete("/keys/{id}", s.handleAPIDeleteKey)
|
||||
|
||||
// Providers
|
||||
r.Get("/providers", s.handleAPIListProviders)
|
||||
r.Get("/providers/{name}", s.handleAPIGetProvider)
|
||||
|
||||
// Scan
|
||||
r.Post("/scan", s.handleAPIScan)
|
||||
r.Get("/scan/progress", s.handleSSEScanProgress)
|
||||
|
||||
// Recon
|
||||
r.Post("/recon", s.handleAPIRecon)
|
||||
r.Get("/recon/progress", s.handleSSEReconProgress)
|
||||
|
||||
// Dorks
|
||||
r.Get("/dorks", s.handleAPIListDorks)
|
||||
r.Post("/dorks", s.handleAPIAddDork)
|
||||
|
||||
// Config
|
||||
r.Get("/config", s.handleAPIGetConfig)
|
||||
r.Put("/config", s.handleAPIUpdateConfig)
|
||||
})
|
||||
}
|
||||
|
||||
// --- Stats ---
|
||||
|
||||
func (s *Server) handleAPIStats(w http.ResponseWriter, r *http.Request) {
|
||||
var totalKeys int
|
||||
if s.cfg.DB != nil {
|
||||
row := s.cfg.DB.SQL().QueryRow("SELECT COUNT(*) FROM findings")
|
||||
_ = row.Scan(&totalKeys)
|
||||
}
|
||||
|
||||
totalProviders := 0
|
||||
if s.cfg.Providers != nil {
|
||||
totalProviders = len(s.cfg.Providers.List())
|
||||
}
|
||||
|
||||
reconSources := 0
|
||||
if s.cfg.ReconEngine != nil {
|
||||
reconSources = len(s.cfg.ReconEngine.List())
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"totalKeys": totalKeys,
|
||||
"totalProviders": totalProviders,
|
||||
"reconSources": reconSources,
|
||||
})
|
||||
}
|
||||
|
||||
// --- Keys ---
|
||||
|
||||
func (s *Server) handleAPIListKeys(w http.ResponseWriter, r *http.Request) {
|
||||
if s.cfg.DB == nil {
|
||||
writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
|
||||
q := r.URL.Query()
|
||||
f := storage.Filters{
|
||||
Provider: q.Get("provider"),
|
||||
Limit: intParam(q.Get("limit"), 50),
|
||||
Offset: intParam(q.Get("offset"), 0),
|
||||
}
|
||||
|
||||
if v := q.Get("verified"); v != "" {
|
||||
b := v == "true" || v == "1"
|
||||
f.Verified = &b
|
||||
}
|
||||
|
||||
findings, err := s.cfg.DB.ListFindingsFiltered(s.cfg.EncKey, f)
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
keys := make([]apiKey, 0, len(findings))
|
||||
for _, f := range findings {
|
||||
keys = append(keys, apiKey{
|
||||
ID: f.ID,
|
||||
ScanID: f.ScanID,
|
||||
ProviderName: f.ProviderName,
|
||||
KeyValue: "", // always masked
|
||||
KeyMasked: f.KeyMasked,
|
||||
Confidence: f.Confidence,
|
||||
SourcePath: f.SourcePath,
|
||||
SourceType: f.SourceType,
|
||||
LineNumber: f.LineNumber,
|
||||
Verified: f.Verified,
|
||||
VerifyStatus: f.VerifyStatus,
|
||||
VerifyHTTPCode: f.VerifyHTTPCode,
|
||||
VerifyMetadata: f.VerifyMetadata,
|
||||
CreatedAt: f.CreatedAt.Format("2006-01-02T15:04:05Z"),
|
||||
})
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusOK, keys)
|
||||
}
|
||||
|
||||
func (s *Server) handleAPIGetKey(w http.ResponseWriter, r *http.Request) {
|
||||
if s.cfg.DB == nil {
|
||||
writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
|
||||
id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid id"})
|
||||
return
|
||||
}
|
||||
|
||||
finding, err := s.cfg.DB.GetFinding(id, s.cfg.EncKey)
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
writeJSON(w, http.StatusNotFound, map[string]string{"error": "not found"})
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// Return masked via apiKey struct
|
||||
writeJSON(w, http.StatusOK, apiKey{
|
||||
ID: finding.ID,
|
||||
ScanID: finding.ScanID,
|
||||
ProviderName: finding.ProviderName,
|
||||
KeyValue: "", // always masked
|
||||
KeyMasked: finding.KeyMasked,
|
||||
Confidence: finding.Confidence,
|
||||
SourcePath: finding.SourcePath,
|
||||
SourceType: finding.SourceType,
|
||||
LineNumber: finding.LineNumber,
|
||||
Verified: finding.Verified,
|
||||
VerifyStatus: finding.VerifyStatus,
|
||||
VerifyHTTPCode: finding.VerifyHTTPCode,
|
||||
VerifyMetadata: finding.VerifyMetadata,
|
||||
CreatedAt: finding.CreatedAt.Format("2006-01-02T15:04:05Z"),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) handleAPIDeleteKey(w http.ResponseWriter, r *http.Request) {
|
||||
if s.cfg.DB == nil {
|
||||
writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
|
||||
id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid id"})
|
||||
return
|
||||
}
|
||||
|
||||
n, err := s.cfg.DB.DeleteFinding(id)
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if n == 0 {
|
||||
writeJSON(w, http.StatusNotFound, map[string]string{"error": "not found"})
|
||||
return
|
||||
}
|
||||
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
|
||||
// --- Providers ---
|
||||
|
||||
// apiProvider is the JSON-friendly representation of a provider.
|
||||
type apiProvider struct {
|
||||
Name string `json:"name"`
|
||||
DisplayName string `json:"displayName"`
|
||||
Tier int `json:"tier"`
|
||||
LastVerified string `json:"lastVerified,omitempty"`
|
||||
Keywords []string `json:"keywords"`
|
||||
}
|
||||
|
||||
func toAPIProvider(p providers.Provider) apiProvider {
|
||||
return apiProvider{
|
||||
Name: p.Name,
|
||||
DisplayName: p.DisplayName,
|
||||
Tier: p.Tier,
|
||||
LastVerified: p.LastVerified,
|
||||
Keywords: p.Keywords,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) handleAPIListProviders(w http.ResponseWriter, r *http.Request) {
|
||||
if s.cfg.Providers == nil {
|
||||
writeJSON(w, http.StatusOK, []interface{}{})
|
||||
return
|
||||
}
|
||||
list := s.cfg.Providers.List()
|
||||
out := make([]apiProvider, len(list))
|
||||
for i, p := range list {
|
||||
out[i] = toAPIProvider(p)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
func (s *Server) handleAPIGetProvider(w http.ResponseWriter, r *http.Request) {
|
||||
if s.cfg.Providers == nil {
|
||||
writeJSON(w, http.StatusNotFound, map[string]string{"error": "not found"})
|
||||
return
|
||||
}
|
||||
|
||||
name := chi.URLParam(r, "name")
|
||||
p, ok := s.cfg.Providers.Get(name)
|
||||
if !ok {
|
||||
writeJSON(w, http.StatusNotFound, map[string]string{"error": "not found"})
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, toAPIProvider(p))
|
||||
}
|
||||
|
||||
// --- Scan ---
|
||||
|
||||
type scanRequest struct {
|
||||
Path string `json:"path"`
|
||||
Verify bool `json:"verify"`
|
||||
Workers int `json:"workers"`
|
||||
}
|
||||
|
||||
func (s *Server) handleAPIScan(w http.ResponseWriter, r *http.Request) {
|
||||
var req scanRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// Launch scan in background -- actual scan wiring happens when ScanEngine
|
||||
// is available. For now, return 202 to indicate the request was accepted.
|
||||
if s.cfg.ScanEngine != nil {
|
||||
go func() {
|
||||
// Background scan execution with SSE progress broadcasting.
|
||||
// Full wiring deferred to serve command integration.
|
||||
s.sse.Broadcast(SSEEvent{Type: "scan:started", Data: map[string]string{
|
||||
"path": req.Path,
|
||||
}})
|
||||
}()
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusAccepted, map[string]string{
|
||||
"status": "started",
|
||||
"message": "scan initiated",
|
||||
})
|
||||
}
|
||||
|
||||
// --- Recon ---
|
||||
|
||||
type reconRequest struct {
|
||||
Query string `json:"query"`
|
||||
Sources []string `json:"sources"`
|
||||
Stealth bool `json:"stealth"`
|
||||
}
|
||||
|
||||
func (s *Server) handleAPIRecon(w http.ResponseWriter, r *http.Request) {
|
||||
var req reconRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if s.cfg.ReconEngine != nil {
|
||||
go func() {
|
||||
cfg := recon.Config{
|
||||
Query: req.Query,
|
||||
EnabledSources: req.Sources,
|
||||
Stealth: req.Stealth,
|
||||
}
|
||||
s.sse.Broadcast(SSEEvent{Type: "recon:started", Data: map[string]string{
|
||||
"query": req.Query,
|
||||
}})
|
||||
findings, err := s.cfg.ReconEngine.SweepAll(context.Background(), cfg)
|
||||
if err != nil {
|
||||
s.sse.Broadcast(SSEEvent{Type: "recon:error", Data: map[string]string{
|
||||
"error": err.Error(),
|
||||
}})
|
||||
return
|
||||
}
|
||||
for _, f := range findings {
|
||||
s.sse.Broadcast(SSEEvent{Type: "recon:finding", Data: f})
|
||||
}
|
||||
s.sse.Broadcast(SSEEvent{Type: "recon:complete", Data: map[string]int{
|
||||
"total": len(findings),
|
||||
}})
|
||||
}()
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusAccepted, map[string]string{
|
||||
"status": "started",
|
||||
"message": "recon initiated",
|
||||
})
|
||||
}
|
||||
|
||||
// --- Dorks ---
|
||||
|
||||
// apiDork is the JSON-friendly representation of a dork.
|
||||
type apiDork struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Source string `json:"source"`
|
||||
Category string `json:"category"`
|
||||
Query string `json:"query"`
|
||||
Description string `json:"description"`
|
||||
Tags []string `json:"tags,omitempty"`
|
||||
}
|
||||
|
||||
func toAPIDork(d dorks.Dork) apiDork {
|
||||
return apiDork{
|
||||
ID: d.ID,
|
||||
Name: d.Name,
|
||||
Source: d.Source,
|
||||
Category: d.Category,
|
||||
Query: d.Query,
|
||||
Description: d.Description,
|
||||
Tags: d.Tags,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) handleAPIListDorks(w http.ResponseWriter, r *http.Request) {
|
||||
source := r.URL.Query().Get("source")
|
||||
|
||||
var list []dorks.Dork
|
||||
if source != "" && s.cfg.Dorks != nil {
|
||||
list = s.cfg.Dorks.ListBySource(source)
|
||||
} else if s.cfg.Dorks != nil {
|
||||
list = s.cfg.Dorks.List()
|
||||
}
|
||||
|
||||
out := make([]apiDork, len(list))
|
||||
for i, d := range list {
|
||||
out[i] = toAPIDork(d)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
type addDorkRequest struct {
|
||||
DorkID string `json:"dorkId"`
|
||||
Name string `json:"name"`
|
||||
Source string `json:"source"`
|
||||
Category string `json:"category"`
|
||||
Query string `json:"query"`
|
||||
Description string `json:"description"`
|
||||
Tags []string `json:"tags"`
|
||||
}
|
||||
|
||||
func (s *Server) handleAPIAddDork(w http.ResponseWriter, r *http.Request) {
|
||||
if s.cfg.DB == nil {
|
||||
writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "database not available"})
|
||||
return
|
||||
}
|
||||
|
||||
var req addDorkRequest
|
||||
if err := readJSON(r, &req); err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
id, err := s.cfg.DB.SaveCustomDork(storage.CustomDork{
|
||||
DorkID: req.DorkID,
|
||||
Name: req.Name,
|
||||
Source: req.Source,
|
||||
Category: req.Category,
|
||||
Query: req.Query,
|
||||
Description: req.Description,
|
||||
Tags: req.Tags,
|
||||
})
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusCreated, map[string]interface{}{"id": id})
|
||||
}
|
||||
|
||||
// --- Config ---
|
||||
|
||||
func (s *Server) handleAPIGetConfig(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusOK, viper.AllSettings())
|
||||
}
|
||||
|
||||
func (s *Server) handleAPIUpdateConfig(w http.ResponseWriter, r *http.Request) {
|
||||
var settings map[string]interface{}
|
||||
if err := readJSON(r, &settings); err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
for k, v := range settings {
|
||||
viper.Set(k, v)
|
||||
}
|
||||
|
||||
// Attempt to persist; ignore error if no config file is set.
|
||||
_ = viper.WriteConfig()
|
||||
|
||||
writeJSON(w, http.StatusOK, map[string]string{"status": "updated"})
|
||||
}
|
||||
|
||||
// --- Helpers ---
|
||||
|
||||
// writeJSON marshals v to JSON and writes it with the given HTTP status code.
|
||||
func writeJSON(w http.ResponseWriter, status int, v interface{}) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
if err := json.NewEncoder(w).Encode(v); err != nil {
|
||||
// Best effort — headers already sent.
|
||||
fmt.Fprintf(w, `{"error":"encode: %s"}`, err)
|
||||
}
|
||||
}
|
||||
|
||||
// readJSON decodes the request body into v.
|
||||
func readJSON(r *http.Request, v interface{}) error {
|
||||
if r.Body == nil {
|
||||
return fmt.Errorf("empty request body")
|
||||
}
|
||||
defer r.Body.Close()
|
||||
return json.NewDecoder(r.Body).Decode(v)
|
||||
}
|
||||
|
||||
// intParam parses a query param as int, returning defaultVal on empty or error.
|
||||
func intParam(s string, defaultVal int) int {
|
||||
if s == "" {
|
||||
return defaultVal
|
||||
}
|
||||
v, err := strconv.Atoi(s)
|
||||
if err != nil || v < 0 {
|
||||
return defaultVal
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// itoa is a small helper for int64 to string conversion.
|
||||
func itoa(v int64) string {
|
||||
return strconv.FormatInt(v, 10)
|
||||
}
|
||||
329
pkg/web/api_test.go
Normal file
329
pkg/web/api_test.go
Normal file
@@ -0,0 +1,329 @@
|
||||
package web
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/salvacybersec/keyhunter/pkg/dorks"
|
||||
"github.com/salvacybersec/keyhunter/pkg/providers"
|
||||
"github.com/salvacybersec/keyhunter/pkg/recon"
|
||||
"github.com/salvacybersec/keyhunter/pkg/storage"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// testServer creates a Server with in-memory DB and minimal registries for testing.
|
||||
func testServer(t *testing.T) (*Server, []byte) {
|
||||
t.Helper()
|
||||
db, err := storage.Open(":memory:")
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { db.Close() })
|
||||
|
||||
encKey := []byte("0123456789abcdef0123456789abcdef") // 32-byte test key
|
||||
|
||||
provReg := providers.NewRegistryFromProviders([]providers.Provider{
|
||||
{Name: "openai", DisplayName: "OpenAI", Tier: 1, Keywords: []string{"sk-"}},
|
||||
{Name: "anthropic", DisplayName: "Anthropic", Tier: 1, Keywords: []string{"sk-ant-"}},
|
||||
})
|
||||
|
||||
dorkReg := dorks.NewRegistryFromDorks([]dorks.Dork{
|
||||
{ID: "gh-openai-1", Name: "OpenAI GitHub", Source: "github", Category: "frontier", Query: "sk-proj- in:file", Description: "Find OpenAI keys on GitHub"},
|
||||
})
|
||||
|
||||
reconEng := recon.NewEngine()
|
||||
|
||||
s := NewServer(ServerConfig{
|
||||
DB: db,
|
||||
EncKey: encKey,
|
||||
Providers: provReg,
|
||||
Dorks: dorkReg,
|
||||
ReconEngine: reconEng,
|
||||
})
|
||||
|
||||
return s, encKey
|
||||
}
|
||||
|
||||
// seedFinding inserts a test finding and returns its ID.
|
||||
func seedFinding(t *testing.T, db *storage.DB, encKey []byte, provider string) int64 {
|
||||
t.Helper()
|
||||
id, err := db.SaveFinding(storage.Finding{
|
||||
ProviderName: provider,
|
||||
KeyValue: "sk-test1234567890abcdefghijklmnop",
|
||||
KeyMasked: "sk-test1...mnop",
|
||||
Confidence: "high",
|
||||
SourcePath: "/tmp/test.py",
|
||||
SourceType: "file",
|
||||
LineNumber: 42,
|
||||
}, encKey)
|
||||
require.NoError(t, err)
|
||||
return id
|
||||
}
|
||||
|
||||
func TestAPIStats(t *testing.T) {
|
||||
s, encKey := testServer(t)
|
||||
seedFinding(t, s.cfg.DB, encKey, "openai")
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/stats", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
assert.Equal(t, "application/json", w.Header().Get("Content-Type"))
|
||||
|
||||
var body map[string]interface{}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
assert.Contains(t, body, "totalKeys")
|
||||
assert.Contains(t, body, "totalProviders")
|
||||
assert.Contains(t, body, "reconSources")
|
||||
}
|
||||
|
||||
func TestAPIListKeys(t *testing.T) {
|
||||
s, encKey := testServer(t)
|
||||
seedFinding(t, s.cfg.DB, encKey, "openai")
|
||||
seedFinding(t, s.cfg.DB, encKey, "anthropic")
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/keys", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
var keys []map[string]interface{}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &keys))
|
||||
assert.Len(t, keys, 2)
|
||||
|
||||
// Keys should be masked (no raw key value exposed)
|
||||
for _, k := range keys {
|
||||
val, ok := k["keyValue"]
|
||||
assert.True(t, ok)
|
||||
assert.Equal(t, "", val, "API must not expose raw key values")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIListKeysFilterByProvider(t *testing.T) {
|
||||
s, encKey := testServer(t)
|
||||
seedFinding(t, s.cfg.DB, encKey, "openai")
|
||||
seedFinding(t, s.cfg.DB, encKey, "anthropic")
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/keys?provider=openai", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
var keys []map[string]interface{}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &keys))
|
||||
assert.Len(t, keys, 1)
|
||||
}
|
||||
|
||||
func TestAPIGetKey(t *testing.T) {
|
||||
s, encKey := testServer(t)
|
||||
id := seedFinding(t, s.cfg.DB, encKey, "openai")
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/keys/"+itoa(id), nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
var body map[string]interface{}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
assert.Equal(t, "openai", body["providerName"])
|
||||
}
|
||||
|
||||
func TestAPIGetKeyNotFound(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/keys/99999", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusNotFound, w.Code)
|
||||
}
|
||||
|
||||
func TestAPIDeleteKey(t *testing.T) {
|
||||
s, encKey := testServer(t)
|
||||
id := seedFinding(t, s.cfg.DB, encKey, "openai")
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodDelete, "/api/v1/keys/"+itoa(id), nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusNoContent, w.Code)
|
||||
}
|
||||
|
||||
func TestAPIDeleteKeyNotFound(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodDelete, "/api/v1/keys/99999", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusNotFound, w.Code)
|
||||
}
|
||||
|
||||
func TestAPIListProviders(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/providers", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
var provs []map[string]interface{}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &provs))
|
||||
assert.Len(t, provs, 2)
|
||||
}
|
||||
|
||||
func TestAPIGetProvider(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/providers/openai", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
var body map[string]interface{}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
|
||||
assert.Equal(t, "openai", body["name"])
|
||||
}
|
||||
|
||||
func TestAPIGetProviderNotFound(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/providers/nonexistent", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusNotFound, w.Code)
|
||||
}
|
||||
|
||||
func TestAPIScan(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
body := `{"path":"/tmp/test","verify":false,"workers":2}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/scan", bytes.NewBufferString(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusAccepted, w.Code)
|
||||
var resp map[string]interface{}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
|
||||
assert.Equal(t, "started", resp["status"])
|
||||
}
|
||||
|
||||
func TestAPIRecon(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
body := `{"query":"openai","sources":["github"],"stealth":false}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/recon", bytes.NewBufferString(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusAccepted, w.Code)
|
||||
var resp map[string]interface{}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
|
||||
assert.Equal(t, "started", resp["status"])
|
||||
}
|
||||
|
||||
func TestAPIListDorks(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/dorks", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
var d []map[string]interface{}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &d))
|
||||
assert.Len(t, d, 1)
|
||||
}
|
||||
|
||||
func TestAPIAddDork(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
body := `{"dorkId":"custom-1","name":"Custom Dork","source":"github","category":"custom","query":"custom query","description":"test"}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/dorks", bytes.NewBufferString(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusCreated, w.Code)
|
||||
var resp map[string]interface{}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
|
||||
assert.Contains(t, resp, "id")
|
||||
}
|
||||
|
||||
func TestAPIGetConfig(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/config", nil)
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
assert.Equal(t, "application/json", w.Header().Get("Content-Type"))
|
||||
}
|
||||
|
||||
func TestAPIUpdateConfig(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
body := `{"scan.workers":"8"}`
|
||||
req := httptest.NewRequest(http.MethodPut, "/api/v1/config", bytes.NewBufferString(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
assert.Equal(t, http.StatusOK, w.Code)
|
||||
}
|
||||
@@ -1,84 +1,34 @@
|
||||
// Package web implements the KeyHunter embedded web dashboard and REST API.
|
||||
package web
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"html/template"
|
||||
"io/fs"
|
||||
"net/http"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/go-chi/chi/v5/middleware"
|
||||
|
||||
"github.com/salvacybersec/keyhunter/pkg/dorks"
|
||||
"github.com/salvacybersec/keyhunter/pkg/engine"
|
||||
"github.com/salvacybersec/keyhunter/pkg/providers"
|
||||
"github.com/salvacybersec/keyhunter/pkg/recon"
|
||||
"github.com/salvacybersec/keyhunter/pkg/storage"
|
||||
)
|
||||
|
||||
// Config holds all dependencies and settings needed by the web server.
|
||||
type Config struct {
|
||||
// ServerConfig holds all dependencies injected into the web Server.
|
||||
type ServerConfig struct {
|
||||
DB *storage.DB
|
||||
EncKey []byte
|
||||
Providers *providers.Registry
|
||||
Dorks *dorks.Registry
|
||||
ScanEngine *engine.Engine
|
||||
ReconEngine *recon.Engine
|
||||
Port int
|
||||
AuthUser string
|
||||
AuthPass string
|
||||
AuthToken string
|
||||
}
|
||||
|
||||
// Server is the KeyHunter web dashboard backed by chi.
|
||||
// Server is the central HTTP server holding all handler dependencies.
|
||||
type Server struct {
|
||||
router chi.Router
|
||||
cfg Config
|
||||
tmpl *template.Template
|
||||
cfg ServerConfig
|
||||
sse *SSEHub
|
||||
}
|
||||
|
||||
// NewServer creates a Server, parsing embedded templates and building routes.
|
||||
func NewServer(cfg Config) (*Server, error) {
|
||||
// Parse all templates from the embedded filesystem.
|
||||
tmpl, err := template.ParseFS(templateFiles, "templates/*.html")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parsing templates: %w", err)
|
||||
// NewServer creates a Server with the given configuration.
|
||||
func NewServer(cfg ServerConfig) *Server {
|
||||
return &Server{
|
||||
cfg: cfg,
|
||||
sse: NewSSEHub(),
|
||||
}
|
||||
|
||||
s := &Server{
|
||||
cfg: cfg,
|
||||
tmpl: tmpl,
|
||||
}
|
||||
|
||||
r := chi.NewRouter()
|
||||
|
||||
// Middleware stack.
|
||||
r.Use(middleware.RealIP)
|
||||
r.Use(middleware.Logger)
|
||||
r.Use(middleware.Recoverer)
|
||||
|
||||
// Auth middleware (no-op when auth fields are empty).
|
||||
r.Use(AuthMiddleware(cfg.AuthUser, cfg.AuthPass, cfg.AuthToken))
|
||||
|
||||
// Static file serving from embedded FS.
|
||||
staticSub, err := fs.Sub(staticFiles, "static")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("creating static sub-filesystem: %w", err)
|
||||
}
|
||||
r.Handle("/static/*", http.StripPrefix("/static/", http.FileServer(http.FS(staticSub))))
|
||||
|
||||
// Page routes.
|
||||
r.Get("/", s.handleOverview)
|
||||
|
||||
s.router = r
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// Router returns the chi router for testing.
|
||||
func (s *Server) Router() chi.Router {
|
||||
return s.router
|
||||
}
|
||||
|
||||
// ListenAndServe starts the HTTP server on the configured port.
|
||||
func (s *Server) ListenAndServe() error {
|
||||
addr := fmt.Sprintf(":%d", s.cfg.Port)
|
||||
return http.ListenAndServe(addr, s.router)
|
||||
}
|
||||
|
||||
115
pkg/web/sse.go
Normal file
115
pkg/web/sse.go
Normal file
@@ -0,0 +1,115 @@
|
||||
package web
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// SSEEvent represents a server-sent event with a type and JSON-serializable data.
|
||||
type SSEEvent struct {
|
||||
Type string `json:"type"`
|
||||
Data interface{} `json:"data"`
|
||||
}
|
||||
|
||||
// SSEHub manages SSE client subscriptions and broadcasts events to all
|
||||
// connected clients. It is safe for concurrent use.
|
||||
type SSEHub struct {
|
||||
mu sync.RWMutex
|
||||
clients map[chan SSEEvent]struct{}
|
||||
}
|
||||
|
||||
// NewSSEHub creates an empty SSE hub ready to accept subscriptions.
|
||||
func NewSSEHub() *SSEHub {
|
||||
return &SSEHub{
|
||||
clients: make(map[chan SSEEvent]struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Subscribe creates a new buffered channel for a client and registers it.
|
||||
// The caller must call Unsubscribe when done.
|
||||
func (h *SSEHub) Subscribe() chan SSEEvent {
|
||||
ch := make(chan SSEEvent, 32)
|
||||
h.mu.Lock()
|
||||
h.clients[ch] = struct{}{}
|
||||
h.mu.Unlock()
|
||||
return ch
|
||||
}
|
||||
|
||||
// Unsubscribe removes a client channel from the hub and closes it.
|
||||
func (h *SSEHub) Unsubscribe(ch chan SSEEvent) {
|
||||
h.mu.Lock()
|
||||
delete(h.clients, ch)
|
||||
h.mu.Unlock()
|
||||
close(ch)
|
||||
}
|
||||
|
||||
// Broadcast sends an event to all connected clients. If a client's buffer is
|
||||
// full the event is dropped for that client (non-blocking send).
|
||||
func (h *SSEHub) Broadcast(evt SSEEvent) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
for ch := range h.clients {
|
||||
select {
|
||||
case ch <- evt:
|
||||
default:
|
||||
// client buffer full, drop event
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ClientCount returns the number of currently connected SSE clients.
|
||||
func (h *SSEHub) ClientCount() int {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
return len(h.clients)
|
||||
}
|
||||
|
||||
// handleSSEScanProgress streams scan progress events to the client via SSE.
|
||||
func (s *Server) handleSSEScanProgress(w http.ResponseWriter, r *http.Request) {
|
||||
s.serveSSE(w, r)
|
||||
}
|
||||
|
||||
// handleSSEReconProgress streams recon progress events to the client via SSE.
|
||||
func (s *Server) handleSSEReconProgress(w http.ResponseWriter, r *http.Request) {
|
||||
s.serveSSE(w, r)
|
||||
}
|
||||
|
||||
// serveSSE is the shared SSE handler for both scan and recon progress endpoints.
|
||||
func (s *Server) serveSSE(w http.ResponseWriter, r *http.Request) {
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
http.Error(w, "streaming unsupported", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
w.Header().Set("X-Accel-Buffering", "no")
|
||||
|
||||
ch := s.sse.Subscribe()
|
||||
defer s.sse.Unsubscribe(ch)
|
||||
|
||||
// Send initial connection event
|
||||
fmt.Fprintf(w, "event: connected\ndata: {}\n\n")
|
||||
flusher.Flush()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
case evt, ok := <-ch:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
data, err := json.Marshal(evt.Data)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
fmt.Fprintf(w, "event: %s\ndata: %s\n\n", evt.Type, data)
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
217
pkg/web/sse_test.go
Normal file
217
pkg/web/sse_test.go
Normal file
@@ -0,0 +1,217 @@
|
||||
package web
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestSSEHubSubscribeUnsubscribe(t *testing.T) {
|
||||
hub := NewSSEHub()
|
||||
|
||||
ch1 := hub.Subscribe()
|
||||
ch2 := hub.Subscribe()
|
||||
assert.Equal(t, 2, hub.ClientCount())
|
||||
|
||||
hub.Unsubscribe(ch1)
|
||||
assert.Equal(t, 1, hub.ClientCount())
|
||||
|
||||
hub.Unsubscribe(ch2)
|
||||
assert.Equal(t, 0, hub.ClientCount())
|
||||
}
|
||||
|
||||
func TestSSEHubBroadcast(t *testing.T) {
|
||||
hub := NewSSEHub()
|
||||
|
||||
ch1 := hub.Subscribe()
|
||||
ch2 := hub.Subscribe()
|
||||
defer hub.Unsubscribe(ch1)
|
||||
defer hub.Unsubscribe(ch2)
|
||||
|
||||
evt := SSEEvent{Type: "scan:progress", Data: map[string]int{"percent": 50}}
|
||||
hub.Broadcast(evt)
|
||||
|
||||
// Both clients should receive the event
|
||||
select {
|
||||
case got := <-ch1:
|
||||
assert.Equal(t, "scan:progress", got.Type)
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("ch1 did not receive event")
|
||||
}
|
||||
|
||||
select {
|
||||
case got := <-ch2:
|
||||
assert.Equal(t, "scan:progress", got.Type)
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("ch2 did not receive event")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSSEHubBroadcastDropsWhenFull(t *testing.T) {
|
||||
hub := NewSSEHub()
|
||||
ch := hub.Subscribe()
|
||||
defer hub.Unsubscribe(ch)
|
||||
|
||||
// Fill the buffer (capacity 32)
|
||||
for i := 0; i < 32; i++ {
|
||||
hub.Broadcast(SSEEvent{Type: "fill", Data: i})
|
||||
}
|
||||
|
||||
// This should NOT block — it drops the event
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
hub.Broadcast(SSEEvent{Type: "overflow", Data: 33})
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
// good, broadcast returned
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Broadcast blocked on full buffer")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSSEHubClientDisconnect(t *testing.T) {
|
||||
hub := NewSSEHub()
|
||||
ch := hub.Subscribe()
|
||||
assert.Equal(t, 1, hub.ClientCount())
|
||||
|
||||
hub.Unsubscribe(ch)
|
||||
assert.Equal(t, 0, hub.ClientCount())
|
||||
|
||||
// Channel should be closed
|
||||
_, ok := <-ch
|
||||
assert.False(t, ok, "channel should be closed after unsubscribe")
|
||||
}
|
||||
|
||||
func TestSSEHTTPHandler(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
// Start the SSE request in a goroutine with a cancelable context
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/scan/progress", nil)
|
||||
req = req.WithContext(ctx)
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
// Run handler in background
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
r.ServeHTTP(w, req)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
// Give handler time to set headers and send initial event
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
// Broadcast an event
|
||||
s.sse.Broadcast(SSEEvent{Type: "scan:finding", Data: map[string]string{"key": "test"}})
|
||||
|
||||
// Give time for event to be written
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
// Cancel the context to disconnect
|
||||
cancel()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("handler did not return after context cancel")
|
||||
}
|
||||
|
||||
// Check response headers
|
||||
assert.Equal(t, "text/event-stream", w.Header().Get("Content-Type"))
|
||||
assert.Equal(t, "no-cache", w.Header().Get("Cache-Control"))
|
||||
|
||||
// Parse SSE events from body
|
||||
body := w.Body.String()
|
||||
assert.Contains(t, body, "event: connected")
|
||||
assert.Contains(t, body, "event: scan:finding")
|
||||
assert.Contains(t, body, "data:")
|
||||
}
|
||||
|
||||
func TestSSEEventFormat(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/recon/progress", nil)
|
||||
req = req.WithContext(ctx)
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
r.ServeHTTP(w, req)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
s.sse.Broadcast(SSEEvent{Type: "recon:complete", Data: map[string]int{"total": 5}})
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
cancel()
|
||||
<-done
|
||||
|
||||
// Verify SSE format: "event: {type}\ndata: {json}\n\n"
|
||||
body := w.Body.String()
|
||||
scanner := bufio.NewScanner(strings.NewReader(body))
|
||||
var foundEvent, foundData bool
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if strings.HasPrefix(line, "event: recon:complete") {
|
||||
foundEvent = true
|
||||
}
|
||||
if strings.HasPrefix(line, "data: ") && strings.Contains(line, `"total"`) {
|
||||
foundData = true
|
||||
}
|
||||
}
|
||||
assert.True(t, foundEvent, "should have event: recon:complete line")
|
||||
assert.True(t, foundData, "should have data line with JSON")
|
||||
}
|
||||
|
||||
func TestSSEClientDisconnectRemovesSubscriber(t *testing.T) {
|
||||
s, _ := testServer(t)
|
||||
|
||||
r := chi.NewRouter()
|
||||
s.mountAPI(r)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/scan/progress", nil)
|
||||
req = req.WithContext(ctx)
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
r.ServeHTTP(w, req)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
assert.Equal(t, 1, s.sse.ClientCount(), "should have 1 subscriber")
|
||||
|
||||
cancel()
|
||||
<-done
|
||||
|
||||
// After disconnect, subscriber should be removed
|
||||
// Give a small moment for cleanup
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
assert.Equal(t, 0, s.sse.ClientCount(), "should have 0 subscribers after disconnect")
|
||||
}
|
||||
Reference in New Issue
Block a user