keyhunter/pkg/storage/findings_test.go

package storage_test

import (
	"database/sql"
	"testing"

	"github.com/salvacybersec/keyhunter/pkg/storage"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	_ "modernc.org/sqlite"
)

func makeTestKey() []byte {
	k := make([]byte, 32)
	for i := range k {
		k[i] = byte(i)
	}
	return k
}

func TestSaveFinding_VerifyFields_RoundTrip(t *testing.T) {
	db, err := storage.Open(":memory:")
	require.NoError(t, err)
	defer db.Close()

	encKey := makeTestKey()

	f := storage.Finding{
		ProviderName:   "openai",
		KeyValue:       "sk-proj-abcdefghijklmnop1234",
		Confidence:     "high",
		SourcePath:     "/tmp/test.env",
		SourceType:     "file",
		LineNumber:     12,
		Verified:       true,
		VerifyStatus:   "live",
		VerifyHTTPCode: 200,
		VerifyMetadata: map[string]string{"org": "Acme", "tier": "plus"},
	}
	_, err = db.SaveFinding(f, encKey)
	require.NoError(t, err)

	out, err := db.ListFindings(encKey)
	require.NoError(t, err)
	require.Len(t, out, 1)

	got := out[0]
	assert.True(t, got.Verified)
	assert.Equal(t, "live", got.VerifyStatus)
	assert.Equal(t, 200, got.VerifyHTTPCode)
	assert.Equal(t, "Acme", got.VerifyMetadata["org"])
	assert.Equal(t, "plus", got.VerifyMetadata["tier"])
	assert.Equal(t, f.KeyValue, got.KeyValue)
}

func TestSaveFinding_VerifyFields_Empty(t *testing.T) {
	db, err := storage.Open(":memory:")
	require.NoError(t, err)
	defer db.Close()

	encKey := makeTestKey()

	f := storage.Finding{
		ProviderName: "anthropic",
		KeyValue:     "sk-ant-api03-abcdefghij1234",
		Confidence:   "medium",
		SourcePath:   "/tmp/cfg.yaml",
		SourceType:   "file",
	}
	_, err = db.SaveFinding(f, encKey)
	require.NoError(t, err)

	out, err := db.ListFindings(encKey)
	require.NoError(t, err)
	require.Len(t, out, 1)

	got := out[0]
	assert.False(t, got.Verified)
	assert.Equal(t, "", got.VerifyStatus)
	assert.Equal(t, 0, got.VerifyHTTPCode)
	assert.Nil(t, got.VerifyMetadata)
}

func TestOpen_MigratesExistingDB(t *testing.T) {
	// Create a temp file-backed DB so we can close and reopen it.
	tmp := t.TempDir() + "/legacy.db"

	// Bootstrap an OLD-schema findings table directly (no verify columns).
	raw, err := sql.Open("sqlite", tmp)
	require.NoError(t, err)
	_, err = raw.Exec(`CREATE TABLE findings (
		id            INTEGER PRIMARY KEY AUTOINCREMENT,
		scan_id       INTEGER,
		provider_name TEXT NOT NULL,
		key_value     BLOB NOT NULL,
		key_masked    TEXT NOT NULL,
		confidence    TEXT NOT NULL,
		source_path   TEXT,
		source_type   TEXT,
		line_number   INTEGER,
		created_at    DATETIME DEFAULT CURRENT_TIMESTAMP
	)`)
	require.NoError(t, err)
	require.NoError(t, raw.Close())

	// Reopen via storage.Open — should migrate the table in-place.
	db, err := storage.Open(tmp)
	require.NoError(t, err)
	defer db.Close()

	rows, err := db.SQL().Query("PRAGMA table_info(findings)")
	require.NoError(t, err)
	defer rows.Close()

	cols := map[string]bool{}
	for rows.Next() {
		var cid int
		var name, ctype string
		var notnull, pk int
		var dflt sql.NullString
		require.NoError(t, rows.Scan(&cid, &name, &ctype, &notnull, &dflt, &pk))
		cols[name] = true
	}
	assert.True(t, cols["verified"], "verified column missing")
	assert.True(t, cols["verify_status"], "verify_status column missing")
	assert.True(t, cols["verify_http_code"], "verify_http_code column missing")
	assert.True(t, cols["verify_metadata_json"], "verify_metadata_json column missing")

	// And verify the round-trip still works after migration.
	encKey := makeTestKey()
	_, err = db.SaveFinding(storage.Finding{
		ProviderName:   "groq",
		KeyValue:       "gsk_abcdefghijklmnop1234",
		Confidence:     "high",
		Verified:       true,
		VerifyStatus:   "live",
		VerifyHTTPCode: 200,
		VerifyMetadata: map[string]string{"model": "mixtral"},
	}, encKey)
	require.NoError(t, err)

	out, err := db.ListFindings(encKey)
	require.NoError(t, err)
	require.Len(t, out, 1)
	assert.Equal(t, "live", out[0].VerifyStatus)
	assert.Equal(t, "mixtral", out[0].VerifyMetadata["model"])
}