keyhunter/pkg/storage/custom_dorks.go

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
}