keyhunter/pkg/dorks/registry_test.go

package dorks_test

import (
	"context"
	"errors"
	"testing"

	"github.com/salvacybersec/keyhunter/pkg/dorks"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func fixture() []dorks.Dork {
	return []dorks.Dork{
		{
			ID:          "openai-github-envfile",
			Name:        "OpenAI API Key in .env files",
			Source:      "github",
			Category:    "frontier",
			Query:       "sk-proj- extension:env",
			Description: "Finds OpenAI project keys exposed in committed .env files",
			Tags:        []string{"openai", "env", "tier1"},
		},
		{
			ID:       "anthropic-github-env",
			Name:     "Anthropic Key in .env",
			Source:   "github",
			Category: "frontier",
			Query:    "sk-ant-api03- extension:env",
		},
		{
			ID:       "shodan-openai-banner",
			Name:     "OpenAI banner on Shodan",
			Source:   "shodan",
			Category: "infrastructure",
			Query:    "product:openai",
		},
	}
}

func TestRegistry_LoadsAndIndexesDorks(t *testing.T) {
	r := dorks.NewRegistryFromDorks(fixture())
	require.NotNil(t, r)
	assert.Len(t, r.List(), 3)
}

func TestRegistry_Get(t *testing.T) {
	r := dorks.NewRegistryFromDorks(fixture())

	d, ok := r.Get("openai-github-envfile")
	require.True(t, ok)
	assert.Equal(t, "github", d.Source)
	assert.Equal(t, "frontier", d.Category)
	assert.Equal(t, "sk-proj- extension:env", d.Query)

	_, ok = r.Get("does-not-exist")
	assert.False(t, ok)
}

func TestRegistry_ListBySource(t *testing.T) {
	r := dorks.NewRegistryFromDorks(fixture())

	gh := r.ListBySource("github")
	assert.Len(t, gh, 2)
	for _, d := range gh {
		assert.Equal(t, "github", d.Source)
	}

	shodan := r.ListBySource("shodan")
	assert.Len(t, shodan, 1)

	assert.Empty(t, r.ListBySource("fofa"))
}

func TestRegistry_ListByCategory(t *testing.T) {
	r := dorks.NewRegistryFromDorks(fixture())

	frontier := r.ListByCategory("frontier")
	assert.Len(t, frontier, 2)
	for _, d := range frontier {
		assert.Equal(t, "frontier", d.Category)
	}

	infra := r.ListByCategory("infrastructure")
	assert.Len(t, infra, 1)

	assert.Empty(t, r.ListByCategory("emerging"))
}

func TestRegistry_Stats(t *testing.T) {
	r := dorks.NewRegistryFromDorks(fixture())

	s := r.Stats()
	assert.Equal(t, 3, s.Total)
	assert.Equal(t, 2, s.BySource["github"])
	assert.Equal(t, 1, s.BySource["shodan"])
	assert.Equal(t, 2, s.ByCategory["frontier"])
	assert.Equal(t, 1, s.ByCategory["infrastructure"])
}

func TestNewRegistry_EmptyDefinitionsTreeOK(t *testing.T) {
	// The embedded definitions tree contains only .gitkeep in this plan.
	// NewRegistry must tolerate that and return an empty but usable Registry.
	r, err := dorks.NewRegistry()
	require.NoError(t, err)
	require.NotNil(t, r)
	assert.GreaterOrEqual(t, len(r.List()), 0)
}

func TestDork_Validate(t *testing.T) {
	tests := []struct {
		name    string
		dork    dorks.Dork
		wantErr bool
	}{
		{
			name:    "valid",
			dork:    dorks.Dork{ID: "x", Source: "github", Query: "foo"},
			wantErr: false,
		},
		{
			name:    "missing id",
			dork:    dorks.Dork{Source: "github", Query: "foo"},
			wantErr: true,
		},
		{
			name:    "missing source",
			dork:    dorks.Dork{ID: "x", Query: "foo"},
			wantErr: true,
		},
		{
			name:    "missing query",
			dork:    dorks.Dork{ID: "x", Source: "github"},
			wantErr: true,
		},
		{
			name:    "unknown source",
			dork:    dorks.Dork{ID: "x", Source: "reddit", Query: "foo"},
			wantErr: true,
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			err := tc.dork.Validate()
			if tc.wantErr {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

func TestRunner_UnknownSourceReturnsErrSourceNotImplemented(t *testing.T) {
	runner := dorks.NewRunner()
	d := dorks.Dork{ID: "x", Source: "shodan", Category: "infrastructure", Query: "product:openai"}

	_, err := runner.Run(context.Background(), d, 10)
	require.Error(t, err)
	assert.True(t, errors.Is(err, dorks.ErrSourceNotImplemented))
}

// stubExecutor is a test-only Executor that records invocations.
type stubExecutor struct {
	source string
	calls  int
}

func (s *stubExecutor) Source() string { return s.source }
func (s *stubExecutor) Execute(_ context.Context, d dorks.Dork, _ int) ([]dorks.Match, error) {
	s.calls++
	return []dorks.Match{{DorkID: d.ID, Source: s.source, Snippet: "hit"}}, nil
}

func TestRunner_RegisterAndDispatch(t *testing.T) {
	runner := dorks.NewRunner()
	stub := &stubExecutor{source: "github"}
	runner.Register(stub)

	got, ok := runner.Executor("github")
	require.True(t, ok)
	assert.Equal(t, "github", got.Source())

	d := dorks.Dork{ID: "t1", Source: "github", Query: "sk-"}
	matches, err := runner.Run(context.Background(), d, 5)
	require.NoError(t, err)
	assert.Equal(t, 1, stub.calls)
	require.Len(t, matches, 1)
	assert.Equal(t, "t1", matches[0].DorkID)
}