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:
salvacybersec
2026-04-06 00:16:33 +03:00
parent fd6efbb4c2
commit 01062b88b1
3 changed files with 319 additions and 0 deletions

View 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)
}