package cmd import ( "bytes" "context" "path/filepath" "strings" "testing" "github.com/salvacybersec/keyhunter/pkg/dorks" "github.com/salvacybersec/keyhunter/pkg/storage" "github.com/spf13/viper" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) // setupDorksTest isolates viper state, points database.path at a tempdir // DB (auto-migrated by storage.Open), and returns the DB path plus a // t.Cleanup-registered flag reset. func setupDorksTest(t *testing.T) string { t.Helper() dir := t.TempDir() dbPath := filepath.Join(dir, "dorks.db") viper.Reset() viper.Set("database.path", dbPath) resetDorksFlags() // Touch the DB once so schema migrations run deterministically even // when the subcommand under test doesn't reach the DB code path. db, err := storage.Open(dbPath) require.NoError(t, err) require.NoError(t, db.Close()) t.Cleanup(func() { resetDorksFlags() viper.Reset() }) return dbPath } // execDorks runs rootCmd with the given args and returns captured stdout // (which Cobra wires for both Out and Err when SetOut is used) alongside // the terminal error from Execute. func execDorks(t *testing.T, args ...string) (string, error) { t.Helper() var buf bytes.Buffer rootCmd.SetOut(&buf) rootCmd.SetErr(&buf) rootCmd.SetArgs(args) err := rootCmd.Execute() return buf.String(), err } // fakeExecutor is an in-memory dorks.Executor used to exercise `dorks run` // without touching the real GitHub API. type fakeExecutor struct { source string matches []dorks.Match err error called int } func (f *fakeExecutor) Source() string { return f.source } func (f *fakeExecutor) Execute(ctx context.Context, d dorks.Dork, limit int) ([]dorks.Match, error) { f.called++ if f.err != nil { return nil, f.err } out := make([]dorks.Match, len(f.matches)) for i, m := range f.matches { m.DorkID = d.ID out[i] = m } return out, nil } func TestDorksList_FilterBySourceAndCategory(t *testing.T) { setupDorksTest(t) out, err := execDorks(t, "dorks", "list", "--source=github", "--category=frontier") require.NoError(t, err) assert.Contains(t, out, "openai-github-envfile") assert.Contains(t, out, "anthropic-github-envfile") // google-ai dorks are frontier github too; shodan must be filtered out. assert.NotContains(t, out, "shodan-") } func TestDorksAdd_PersistsCustomDorkAndShowsInList(t *testing.T) { setupDorksTest(t) _, err := execDorks(t, "dorks", "add", "--source=github", "--category=frontier", "--id=my-custom-dork", "--name=Custom", "--query=foo extension:env", "--description=test dork") require.NoError(t, err) resetDorksFlags() out, err := execDorks(t, "dorks", "list", "--source=github") require.NoError(t, err) assert.Contains(t, out, "*my-custom-dork") } func TestDorksAdd_RejectsInvalidSource(t *testing.T) { setupDorksTest(t) _, err := execDorks(t, "dorks", "add", "--source=notreal", "--category=frontier", "--id=x", "--query=foo") require.Error(t, err) assert.Contains(t, err.Error(), "not one of") } func TestDorksAdd_RejectsEmbeddedIDCollision(t *testing.T) { setupDorksTest(t) _, err := execDorks(t, "dorks", "add", "--source=github", "--category=frontier", "--id=openai-github-envfile", "--query=foo") require.Error(t, err) assert.Contains(t, err.Error(), "collides with an embedded dork") } func TestDorksDelete_EmbeddedRefused(t *testing.T) { setupDorksTest(t) _, err := execDorks(t, "dorks", "delete", "openai-github-envfile") require.Error(t, err) assert.Contains(t, err.Error(), "embedded dorks cannot be deleted") } func TestDorksDelete_RemovesCustomDork(t *testing.T) { setupDorksTest(t) _, err := execDorks(t, "dorks", "add", "--source=github", "--category=frontier", "--id=to-delete", "--query=bar") require.NoError(t, err) resetDorksFlags() out, err := execDorks(t, "dorks", "delete", "to-delete") require.NoError(t, err) assert.Contains(t, out, "Deleted custom dork") resetDorksFlags() out, err = execDorks(t, "dorks", "list", "--source=github") require.NoError(t, err) assert.NotContains(t, out, "*to-delete") } func TestDorksRun_ShodanReturnsNotImplemented(t *testing.T) { setupDorksTest(t) // Use a shodan dork id if one exists; if none, add one via add first. // Shodan live execution is not wired in Phase 8. _, err := execDorks(t, "dorks", "add", "--source=shodan", "--category=infrastructure", "--id=shodan-test", "--query=foo") require.NoError(t, err) resetDorksFlags() _, err = execDorks(t, "dorks", "run", "--source=shodan", "--id=shodan-test") require.Error(t, err) assert.Contains(t, err.Error(), "not yet implemented") assert.Contains(t, err.Error(), "Phase 9-16") } func TestDorksRun_GitHubMissingTokenReturnsAuthHint(t *testing.T) { setupDorksTest(t) // No GITHUB_TOKEN set, no dorks.github.token — executor should yield // ErrMissingAuth wrapped with a setup hint. viper.Set("dorks.github.token", "") _, err := execDorks(t, "dorks", "run", "--source=github", "--id=openai-github-envfile", "--limit=1") require.Error(t, err) assert.Contains(t, err.Error(), "GITHUB_TOKEN") } func TestDorksRun_GitHubWithInjectedFakeExecutor(t *testing.T) { setupDorksTest(t) fake := &fakeExecutor{ source: "github", matches: []dorks.Match{ {Source: "github", URL: "https://example.com/x", Path: "repo/x.env", Snippet: "sk-proj-XXXX"}, }, } original := newGitHubExecutor newGitHubExecutor = func() dorks.Executor { return fake } t.Cleanup(func() { newGitHubExecutor = original }) out, err := execDorks(t, "dorks", "run", "--source=github", "--id=openai-github-envfile", "--limit=1") require.NoError(t, err) assert.Equal(t, 1, fake.called) assert.Contains(t, out, "openai-github-envfile") assert.Contains(t, out, "https://example.com/x") } func TestDorksExport_YAMLContainsEmbeddedAndCustom(t *testing.T) { setupDorksTest(t) _, err := execDorks(t, "dorks", "add", "--source=github", "--category=frontier", "--id=export-custom", "--query=baz") require.NoError(t, err) resetDorksFlags() out, err := execDorks(t, "dorks", "export", "--format=yaml") require.NoError(t, err) assert.Contains(t, out, "openai-github-envfile") assert.Contains(t, out, "export-custom") // yaml decoded output should contain `id:` keys. assert.True(t, strings.Contains(out, "id:")) } func TestDorksInfo_EmbeddedAndCustom(t *testing.T) { setupDorksTest(t) out, err := execDorks(t, "dorks", "info", "openai-github-envfile") require.NoError(t, err) assert.Contains(t, out, "Origin: embedded") assert.Contains(t, out, "Query:") resetDorksFlags() _, err = execDorks(t, "dorks", "add", "--source=github", "--category=frontier", "--id=info-custom", "--query=qux") require.NoError(t, err) resetDorksFlags() out, err = execDorks(t, "dorks", "info", "info-custom") require.NoError(t, err) assert.Contains(t, out, "Origin: custom") }