keyhunter/cmd/keys_test.go

package cmd

import (
	"bytes"
	"encoding/csv"
	"encoding/json"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"

	"github.com/salvacybersec/keyhunter/pkg/storage"
	"github.com/spf13/viper"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// seedKeysDB creates a fresh SQLite database populated with three findings
// (two openai, one anthropic; one openai row marked verified) and wires
// viper + env so the keys subcommands pick the same DB up via openDBWithKey.
// A t.Cleanup resets global state so flag/viper bleed between tests is
// impossible.
func seedKeysDB(t *testing.T) string {
	t.Helper()
	dir := t.TempDir()
	dbPath := filepath.Join(dir, "keys.db")

	viper.Reset()
	viper.Set("database.path", dbPath)
	t.Setenv("KEYHUNTER_PASSPHRASE", "test-pass")
	resetKeysFlags()

	db, err := storage.Open(dbPath)
	require.NoError(t, err)
	encKey, err := loadOrCreateEncKey(db, "test-pass")
	require.NoError(t, err)

	seed := []storage.Finding{
		{ProviderName: "openai", KeyValue: "sk-aaaaaaaaaaaaaaaaaaaa", KeyMasked: "sk-aaaa...aaaa", Confidence: "high", SourcePath: "a.go", LineNumber: 10},
		{ProviderName: "openai", KeyValue: "sk-bbbbbbbbbbbbbbbbbbbb", KeyMasked: "sk-bbbb...bbbb", Confidence: "medium", SourcePath: "b.go", LineNumber: 20, Verified: true, VerifyStatus: "live"},
		{ProviderName: "anthropic", KeyValue: "sk-ant-cccccccccccccccc", KeyMasked: "sk-ant-c...cccc", Confidence: "high", SourcePath: "c.go", LineNumber: 30},
	}
	for _, f := range seed {
		_, err := db.SaveFinding(f, encKey)
		require.NoError(t, err)
	}
	require.NoError(t, db.Close())

	t.Cleanup(func() {
		viper.Reset()
		resetKeysFlags()
	})
	return dbPath
}

func resetKeysFlags() {
	flagKeysUnmask = false
	flagKeysProvider = ""
	flagKeysVerified = false
	flagKeysLimit = 0
	flagKeysFormat = "json"
	flagKeysOutFile = ""
	flagKeysYes = false
}

func TestKeysList_Default(t *testing.T) {
	_ = seedKeysDB(t)
	var out bytes.Buffer
	keysListCmd.SetOut(&out)
	keysListCmd.SetErr(&out)
	err := keysListCmd.RunE(keysListCmd, []string{})
	require.NoError(t, err)

	s := out.String()
	assert.Contains(t, s, "openai")
	assert.Contains(t, s, "anthropic")
	assert.Contains(t, s, "sk-aaaa...aaaa")
	assert.Contains(t, s, "sk-ant-c...cccc")
	// masked by default: full plaintext key must not appear
	assert.NotContains(t, s, "sk-aaaaaaaaaaaaaaaaaaaa")
	assert.Contains(t, s, "3 key(s).")
}

func TestKeysList_FilterProvider(t *testing.T) {
	_ = seedKeysDB(t)
	flagKeysProvider = "openai"
	var out bytes.Buffer
	keysListCmd.SetOut(&out)
	keysListCmd.SetErr(&out)
	err := keysListCmd.RunE(keysListCmd, []string{})
	require.NoError(t, err)

	s := out.String()
	assert.Contains(t, s, "openai")
	assert.NotContains(t, s, "anthropic")
	assert.Contains(t, s, "2 key(s).")
}

func TestKeysShow_Hit(t *testing.T) {
	dbPath := seedKeysDB(t)
	id := firstFindingID(t, dbPath)

	var out bytes.Buffer
	keysShowCmd.SetOut(&out)
	keysShowCmd.SetErr(&out)
	err := keysShowCmd.RunE(keysShowCmd, []string{strconv.FormatInt(id, 10)})
	require.NoError(t, err)

	s := out.String()
	// show is always unmasked: full plaintext must be present
	assert.Contains(t, s, "sk-aaaaaaaaaaaaaaaaaaaa")
	assert.Contains(t, s, "Provider:")
	assert.Contains(t, s, "openai")
}

func TestKeysShow_Miss(t *testing.T) {
	_ = seedKeysDB(t)
	var out bytes.Buffer
	keysShowCmd.SetOut(&out)
	keysShowCmd.SetErr(&out)
	err := keysShowCmd.RunE(keysShowCmd, []string{"9999"})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "no finding with id 9999")
}

func TestKeysExport_JSON(t *testing.T) {
	_ = seedKeysDB(t)
	flagKeysFormat = "json"
	flagKeysOutFile = ""

	var out bytes.Buffer
	keysExportCmd.SetOut(&out)
	keysExportCmd.SetErr(&out)
	err := keysExportCmd.RunE(keysExportCmd, []string{})
	require.NoError(t, err)

	var parsed []map[string]interface{}
	require.NoError(t, json.Unmarshal(out.Bytes(), &parsed))
	assert.Len(t, parsed, 3)

	// Exports are unmasked — at least one row should carry the plaintext key.
	found := false
	for _, row := range parsed {
		if row["key"] == "sk-aaaaaaaaaaaaaaaaaaaa" {
			found = true
			break
		}
	}
	assert.True(t, found, "expected unmasked plaintext key in JSON export")
}

func TestKeysExport_CSVFile(t *testing.T) {
	_ = seedKeysDB(t)
	outFile := filepath.Join(t.TempDir(), "export.csv")
	flagKeysFormat = "csv"
	flagKeysOutFile = outFile

	var out bytes.Buffer
	keysExportCmd.SetOut(&out)
	keysExportCmd.SetErr(&out)
	err := keysExportCmd.RunE(keysExportCmd, []string{})
	require.NoError(t, err)

	data, err := os.ReadFile(outFile)
	require.NoError(t, err)

	r := csv.NewReader(strings.NewReader(string(data)))
	rows, err := r.ReadAll()
	require.NoError(t, err)

	require.Len(t, rows, 4) // header + 3 data
	assert.Equal(t, "id", rows[0][0])
	assert.Equal(t, "provider", rows[0][1])
}

func TestKeysDelete_WithYes(t *testing.T) {
	dbPath := seedKeysDB(t)
	id := firstFindingID(t, dbPath)
	flagKeysYes = true

	var out bytes.Buffer
	keysDeleteCmd.SetOut(&out)
	keysDeleteCmd.SetErr(&out)
	err := keysDeleteCmd.RunE(keysDeleteCmd, []string{strconv.FormatInt(id, 10)})
	require.NoError(t, err)
	assert.Contains(t, out.String(), "Deleted finding")

	// Verify list now returns 2.
	flagKeysYes = false
	var out2 bytes.Buffer
	keysListCmd.SetOut(&out2)
	keysListCmd.SetErr(&out2)
	require.NoError(t, keysListCmd.RunE(keysListCmd, []string{}))
	assert.Contains(t, out2.String(), "2 key(s).")
}

// NOTE: TestKeysCopy is omitted — the atotto/clipboard backend is unavailable
// in headless CI environments. The copy command path is exercised manually.
// NOTE: TestKeysVerify is omitted — verify requires live network calls to
// provider endpoints which is out of scope for unit tests.

// firstFindingID returns the smallest id present in the findings table of the
// seeded database. Used so tests do not depend on SQLite rowid sequencing.
func firstFindingID(t *testing.T, dbPath string) int64 {
	t.Helper()
	db, err := storage.Open(dbPath)
	require.NoError(t, err)
	defer db.Close()
	row := db.SQL().QueryRow(`SELECT MIN(id) FROM findings`)
	var id int64
	require.NoError(t, row.Scan(&id))
	return id
}