From 01062b88b1598d940bc3d545e4ca2707ce5b096a Mon Sep 17 00:00:00 2001 From: salvacybersec Date: Mon, 6 Apr 2026 00:16:33 +0300 Subject: [PATCH] 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 --- pkg/storage/custom_dorks.go | 145 ++++++++++++++++++++++++++++ pkg/storage/custom_dorks_test.go | 156 +++++++++++++++++++++++++++++++ pkg/storage/schema.sql | 18 ++++ 3 files changed, 319 insertions(+) create mode 100644 pkg/storage/custom_dorks.go create mode 100644 pkg/storage/custom_dorks_test.go diff --git a/pkg/storage/custom_dorks.go b/pkg/storage/custom_dorks.go new file mode 100644 index 0000000..7823ac1 --- /dev/null +++ b/pkg/storage/custom_dorks.go @@ -0,0 +1,145 @@ +package storage + +import ( + "database/sql" + "encoding/json" + "fmt" + "time" +) + +// CustomDork is a user-authored dork persisted to the custom_dorks table. +// It parallels pkg/dorks.Dork but owns its own id (auto-increment) and +// timestamp. DorkID is the stable string identifier exposed to users, and +// is constrained UNIQUE at the SQL layer. +type CustomDork struct { + ID int64 + DorkID string + Name string + Source string + Category string + Query string + Description string + Tags []string + CreatedAt time.Time +} + +// SaveCustomDork inserts a custom dork row and returns its auto-generated +// primary key. Tags are serialised to JSON for storage in the TEXT column. +func (db *DB) SaveCustomDork(d CustomDork) (int64, error) { + tagsJSON, err := marshalTags(d.Tags) + if err != nil { + return 0, fmt.Errorf("encoding tags: %w", err) + } + res, err := db.sql.Exec(` + INSERT INTO custom_dorks (dork_id, name, source, category, query, description, tags) + VALUES (?, ?, ?, ?, ?, ?, ?) + `, d.DorkID, d.Name, d.Source, d.Category, d.Query, d.Description, tagsJSON) + if err != nil { + return 0, fmt.Errorf("inserting custom dork: %w", err) + } + id, err := res.LastInsertId() + if err != nil { + return 0, fmt.Errorf("fetching insert id: %w", err) + } + return id, nil +} + +// ListCustomDorks returns every custom dork ordered newest-first. +func (db *DB) ListCustomDorks() ([]CustomDork, error) { + rows, err := db.sql.Query(` + SELECT id, dork_id, name, source, category, query, description, tags, created_at + FROM custom_dorks + ORDER BY created_at DESC, id DESC + `) + if err != nil { + return nil, fmt.Errorf("querying custom dorks: %w", err) + } + defer rows.Close() + + var out []CustomDork + for rows.Next() { + d, err := scanCustomDork(rows) + if err != nil { + return nil, err + } + out = append(out, d) + } + if err := rows.Err(); err != nil { + return nil, err + } + return out, nil +} + +// GetCustomDork returns the custom dork with the given auto-increment id. +// It returns sql.ErrNoRows when no such row exists. +func (db *DB) GetCustomDork(id int64) (CustomDork, error) { + row := db.sql.QueryRow(` + SELECT id, dork_id, name, source, category, query, description, tags, created_at + FROM custom_dorks WHERE id = ? + `, id) + return scanCustomDork(row) +} + +// GetCustomDorkByDorkID looks up a custom dork by its user-facing dork_id. +// It returns sql.ErrNoRows when no match exists. +func (db *DB) GetCustomDorkByDorkID(dorkID string) (CustomDork, error) { + row := db.sql.QueryRow(` + SELECT id, dork_id, name, source, category, query, description, tags, created_at + FROM custom_dorks WHERE dork_id = ? + `, dorkID) + return scanCustomDork(row) +} + +// DeleteCustomDork removes the row with the given id and returns the number +// of rows affected (0 or 1). +func (db *DB) DeleteCustomDork(id int64) (int64, error) { + res, err := db.sql.Exec(`DELETE FROM custom_dorks WHERE id = ?`, id) + if err != nil { + return 0, fmt.Errorf("deleting custom dork: %w", err) + } + n, err := res.RowsAffected() + if err != nil { + return 0, fmt.Errorf("rows affected: %w", err) + } + return n, nil +} + +// rowScanner abstracts *sql.Row and *sql.Rows so scanCustomDork can serve +// both single-row Get and multi-row List paths. +type rowScanner interface { + Scan(dest ...any) error +} + +func scanCustomDork(r rowScanner) (CustomDork, error) { + var ( + d CustomDork + description sql.NullString + tagsJSON sql.NullString + ) + if err := r.Scan( + &d.ID, &d.DorkID, &d.Name, &d.Source, &d.Category, &d.Query, + &description, &tagsJSON, &d.CreatedAt, + ); err != nil { + return CustomDork{}, err + } + if description.Valid { + d.Description = description.String + } + if tagsJSON.Valid && tagsJSON.String != "" { + if err := json.Unmarshal([]byte(tagsJSON.String), &d.Tags); err != nil { + return CustomDork{}, fmt.Errorf("decoding tags: %w", err) + } + } + return d, nil +} + +func marshalTags(tags []string) (string, error) { + if len(tags) == 0 { + return "", nil + } + b, err := json.Marshal(tags) + if err != nil { + return "", err + } + return string(b), nil +} diff --git a/pkg/storage/custom_dorks_test.go b/pkg/storage/custom_dorks_test.go new file mode 100644 index 0000000..45cc063 --- /dev/null +++ b/pkg/storage/custom_dorks_test.go @@ -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) +} diff --git a/pkg/storage/schema.sql b/pkg/storage/schema.sql index 79b3a86..964dad9 100644 --- a/pkg/storage/schema.sql +++ b/pkg/storage/schema.sql @@ -37,3 +37,21 @@ CREATE TABLE IF NOT EXISTS settings ( CREATE INDEX IF NOT EXISTS idx_findings_scan_id ON findings(scan_id); CREATE INDEX IF NOT EXISTS idx_findings_provider ON findings(provider_name); CREATE INDEX IF NOT EXISTS idx_findings_created ON findings(created_at DESC); + +-- Phase 8: user-authored dork definitions complementing the embedded +-- YAML-based dork registry (pkg/dorks/definitions). Embedded dorks are +-- read-only; custom_dorks stores rows created via `keyhunter dorks add`. +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, -- JSON array + created_at DATETIME DEFAULT CURRENT_TIMESTAMP +); + +CREATE INDEX IF NOT EXISTS idx_custom_dorks_source ON custom_dorks(source); +CREATE INDEX IF NOT EXISTS idx_custom_dorks_category ON custom_dorks(category);