package sources import ( "context" "net/http" "net/http/httptest" "testing" "time" "github.com/salvacybersec/keyhunter/pkg/providers" "github.com/salvacybersec/keyhunter/pkg/recon" ) func TestVirusTotal_Name(t *testing.T) { s := &VirusTotalSource{} if s.Name() != "virustotal" { t.Fatalf("expected virustotal, got %s", s.Name()) } } func TestVirusTotal_Enabled(t *testing.T) { s := &VirusTotalSource{} if s.Enabled(recon.Config{}) { t.Fatal("VirusTotalSource should be disabled without API key") } s.APIKey = "test-key" if !s.Enabled(recon.Config{}) { t.Fatal("VirusTotalSource should be enabled with API key") } } func TestVirusTotal_Sweep(t *testing.T) { mux := http.NewServeMux() mux.HandleFunc("/intelligence/search", func(w http.ResponseWriter, r *http.Request) { if r.Header.Get("x-apikey") != "test-key" { http.Error(w, "unauthorized", http.StatusUnauthorized) return } w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{ "data": [{ "id": "abc123def456", "attributes": { "meaningful_name": "malware.exe", "tags": ["trojan"], "api_key": "sk-proj-ABCDEF1234567890abcdef" } }] }`)) }) srv := httptest.NewServer(mux) defer srv.Close() reg := providers.NewRegistryFromProviders([]providers.Provider{ {Name: "openai", Keywords: []string{"sk-proj-"}}, }) s := &VirusTotalSource{ APIKey: "test-key", BaseURL: srv.URL, Registry: reg, Client: NewClient(), } out := make(chan recon.Finding, 10) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() err := s.Sweep(ctx, "", out) close(out) if err != nil { t.Fatalf("Sweep error: %v", err) } var findings []recon.Finding for f := range out { findings = append(findings, f) } if len(findings) == 0 { t.Fatal("expected at least one finding from VirusTotal") } if findings[0].SourceType != "recon:virustotal" { t.Fatalf("expected recon:virustotal, got %s", findings[0].SourceType) } } func TestVirusTotal_Sweep_Empty(t *testing.T) { mux := http.NewServeMux() mux.HandleFunc("/intelligence/search", func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"data": []}`)) }) srv := httptest.NewServer(mux) defer srv.Close() reg := providers.NewRegistryFromProviders([]providers.Provider{ {Name: "openai", Keywords: []string{"sk-proj-"}}, }) s := &VirusTotalSource{ APIKey: "test-key", BaseURL: srv.URL, Registry: reg, Client: NewClient(), } out := make(chan recon.Finding, 10) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() err := s.Sweep(ctx, "", out) close(out) if err != nil { t.Fatalf("Sweep error: %v", err) } var findings []recon.Finding for f := range out { findings = append(findings, f) } if len(findings) != 0 { t.Fatalf("expected no findings, got %d", len(findings)) } }