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
145
pkg/storage/custom_dorks.go
Normal file
145
pkg/storage/custom_dorks.go
Normal file
@@ -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
|
||||
}
|
||||
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)
|
||||
}
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user