feat(08-01): add custom_dorks table and CRUD for user-authored dorks
- schema.sql: CREATE TABLE IF NOT EXISTS custom_dorks with unique dork_id, source/category indexes, and tags stored as JSON TEXT - custom_dorks.go: Save/List/Get/GetByDorkID/Delete with JSON tag round-trip - Tests: round-trip, newest-first ordering, not-found, unique constraint, delete no-op, schema migration idempotency
This commit is contained in:
156
pkg/storage/custom_dorks_test.go
Normal file
156
pkg/storage/custom_dorks_test.go
Normal file
@@ -0,0 +1,156 @@
|
||||
package storage_test
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/salvacybersec/keyhunter/pkg/storage"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func sampleCustomDork() storage.CustomDork {
|
||||
return storage.CustomDork{
|
||||
DorkID: "user-openai-envfile",
|
||||
Name: "User OpenAI .env",
|
||||
Source: "github",
|
||||
Category: "frontier",
|
||||
Query: "sk-proj- extension:env",
|
||||
Description: "user-added dork",
|
||||
Tags: []string{"openai", "env"},
|
||||
}
|
||||
}
|
||||
|
||||
func TestSaveCustomDork_RoundTrip(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
id, err := db.SaveCustomDork(sampleCustomDork())
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, id, int64(0))
|
||||
|
||||
got, err := db.GetCustomDork(id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "user-openai-envfile", got.DorkID)
|
||||
assert.Equal(t, "github", got.Source)
|
||||
assert.Equal(t, "frontier", got.Category)
|
||||
assert.Equal(t, "sk-proj- extension:env", got.Query)
|
||||
assert.Equal(t, "user-added dork", got.Description)
|
||||
assert.Equal(t, []string{"openai", "env"}, got.Tags)
|
||||
assert.False(t, got.CreatedAt.IsZero())
|
||||
}
|
||||
|
||||
func TestListCustomDorks_NewestFirst(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
first := sampleCustomDork()
|
||||
first.DorkID = "first"
|
||||
_, err = db.SaveCustomDork(first)
|
||||
require.NoError(t, err)
|
||||
|
||||
second := sampleCustomDork()
|
||||
second.DorkID = "second"
|
||||
_, err = db.SaveCustomDork(second)
|
||||
require.NoError(t, err)
|
||||
|
||||
list, err := db.ListCustomDorks()
|
||||
require.NoError(t, err)
|
||||
require.Len(t, list, 2)
|
||||
assert.Equal(t, "second", list[0].DorkID, "newest should be first")
|
||||
assert.Equal(t, "first", list[1].DorkID)
|
||||
}
|
||||
|
||||
func TestGetCustomDork_NotFound(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
_, err = db.GetCustomDork(9999)
|
||||
require.Error(t, err)
|
||||
assert.True(t, errors.Is(err, sql.ErrNoRows))
|
||||
}
|
||||
|
||||
func TestGetCustomDorkByDorkID(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
_, err = db.SaveCustomDork(sampleCustomDork())
|
||||
require.NoError(t, err)
|
||||
|
||||
got, err := db.GetCustomDorkByDorkID("user-openai-envfile")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "user-openai-envfile", got.DorkID)
|
||||
|
||||
_, err = db.GetCustomDorkByDorkID("missing")
|
||||
require.Error(t, err)
|
||||
assert.True(t, errors.Is(err, sql.ErrNoRows))
|
||||
}
|
||||
|
||||
func TestDeleteCustomDork(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
id, err := db.SaveCustomDork(sampleCustomDork())
|
||||
require.NoError(t, err)
|
||||
|
||||
n, err := db.DeleteCustomDork(id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1), n)
|
||||
|
||||
_, err = db.GetCustomDork(id)
|
||||
assert.True(t, errors.Is(err, sql.ErrNoRows))
|
||||
|
||||
// Deleting again is a no-op (0 rows affected, no error).
|
||||
n, err = db.DeleteCustomDork(id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(0), n)
|
||||
}
|
||||
|
||||
func TestSaveCustomDork_UniqueDorkID(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
_, err = db.SaveCustomDork(sampleCustomDork())
|
||||
require.NoError(t, err)
|
||||
|
||||
// Duplicate dork_id must fail the UNIQUE constraint.
|
||||
_, err = db.SaveCustomDork(sampleCustomDork())
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestSchemaMigration_Idempotent(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
require.NoError(t, err)
|
||||
defer db.Close()
|
||||
|
||||
// Re-running the CREATE TABLE IF NOT EXISTS on the same connection
|
||||
// must be a no-op. We verify by executing CREATE TABLE again manually.
|
||||
_, err = db.SQL().Exec(`
|
||||
CREATE TABLE IF NOT EXISTS custom_dorks (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
dork_id TEXT NOT NULL UNIQUE,
|
||||
name TEXT NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
category TEXT NOT NULL,
|
||||
query TEXT NOT NULL,
|
||||
description TEXT,
|
||||
tags TEXT,
|
||||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
`)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// And that pre-existing rows survive the idempotent re-exec.
|
||||
_, err = db.SaveCustomDork(sampleCustomDork())
|
||||
require.NoError(t, err)
|
||||
list, err := db.ListCustomDorks()
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, list, 1)
|
||||
}
|
||||
Reference in New Issue
Block a user