diff --git a/internal/db/mysql.go b/internal/db/mysql.go index 96c2f391a..59432fab6 100644 --- a/internal/db/mysql.go +++ b/internal/db/mysql.go @@ -151,6 +151,9 @@ func (m *mysqlDatabase) EnsureSchema(ctx context.Context) error { -- Trace trace JSON, + -- CPE (Common Platform Enumeration) + cpe JSON, + INDEX idx_timestamp (timestamp), INDEX idx_url (url(255)), INDEX idx_host (host), @@ -163,9 +166,38 @@ func (m *mysqlDatabase) EnsureSchema(ctx context.Context) error { return fmt.Errorf("failed to create schema: %w", err) } + // Back-compat for databases whose schema was created before CPE support. + // New installs already get this column via the CREATE TABLE above; this + // path only matters for in-place upgrades. + // TODO: replace these ad-hoc ensureColumn calls with a proper migration + // framework (e.g. golang-migrate / goose) once more schema changes accumulate. + if err := m.ensureColumn(ctx, "cpe", "JSON"); err != nil { + return fmt.Errorf("failed to ensure cpe column: %w", err) + } + return nil } +func (m *mysqlDatabase) ensureColumn(ctx context.Context, column, definition string) error { + var count int + err := m.db.QueryRowContext(ctx, + `SELECT COUNT(*) FROM information_schema.columns + WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`, + m.cfg.TableName, column, + ).Scan(&count) + if err != nil { + return err + } + if count > 0 { + return nil + } + _, err = m.db.ExecContext(ctx, + fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s", + quoteIdentifier(m.cfg.TableName), quoteIdentifier(column), definition), + ) + return err +} + func (m *mysqlDatabase) InsertBatch(ctx context.Context, results []runner.Result) error { if len(results) == 0 { return nil @@ -193,7 +225,8 @@ func (m *mysqlDatabase) InsertBatch(ctx context.Context, results []runner.Result words, `+"`lines`"+`, header, extracts, extract_regex, chain, chain_status_codes, headless_body, screenshot_bytes, screenshot_path, screenshot_path_rel, stored_response_path, - knowledgebase, link_request, trace + knowledgebase, link_request, trace, + cpe ) VALUES ( ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, @@ -205,7 +238,8 @@ func (m *mysqlDatabase) InsertBatch(ctx context.Context, results []runner.Result ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, - ?, ?, ? + ?, ?, ?, + ? )`, tableName) stmt, err := tx.PrepareContext(ctx, query) @@ -236,6 +270,7 @@ func (m *mysqlDatabase) InsertBatch(ctx context.Context, results []runner.Result kbJSON, _ := json.Marshal(r.KnowledgeBase) linkReqJSON, _ := json.Marshal(r.LinkRequest) traceJSON, _ := json.Marshal(r.Trace) + cpeJSON, _ := json.Marshal(r.CPE) _, err = stmt.ExecContext(ctx, r.Timestamp, r.URL, r.Input, r.Host, r.Port, r.Scheme, r.Path, r.Method, r.FinalURL, @@ -249,6 +284,7 @@ func (m *mysqlDatabase) InsertBatch(ctx context.Context, results []runner.Result chainJSON, chainStatusJSON, r.HeadlessBody, r.ScreenshotBytes, r.ScreenshotPath, r.ScreenshotPathRel, r.StoredResponsePath, kbJSON, linkReqJSON, traceJSON, + cpeJSON, ) if err != nil { return fmt.Errorf("failed to insert result: %w", err) diff --git a/internal/db/postgres.go b/internal/db/postgres.go index 0eca9831a..24cbb70bf 100644 --- a/internal/db/postgres.go +++ b/internal/db/postgres.go @@ -150,15 +150,27 @@ func (p *postgresDatabase) EnsureSchema(ctx context.Context) error { link_request JSONB, -- Trace - trace JSONB + trace JSONB, + + -- CPE (Common Platform Enumeration) + cpe JSONB ); + -- Back-compat for databases whose schema was created before CPE support. + -- New installs already get this column via the CREATE TABLE above; this + -- statement only matters for in-place upgrades. + -- TODO: replace these ad-hoc ALTER TABLE statements with a proper + -- migration framework (e.g. golang-migrate / goose) once more schema + -- changes accumulate. + ALTER TABLE %s ADD COLUMN IF NOT EXISTS cpe JSONB; + CREATE INDEX IF NOT EXISTS %s ON %s(timestamp DESC); CREATE INDEX IF NOT EXISTS %s ON %s(url); CREATE INDEX IF NOT EXISTS %s ON %s(host); CREATE INDEX IF NOT EXISTS %s ON %s(status_code); CREATE INDEX IF NOT EXISTS %s ON %s USING GIN(tech); `, + tableName, tableName, idxTimestamp, tableName, idxURL, tableName, @@ -201,7 +213,8 @@ func (p *postgresDatabase) InsertBatch(ctx context.Context, results []runner.Res words, lines, header, extracts, extract_regex, chain, chain_status_codes, headless_body, screenshot_bytes, screenshot_path, screenshot_path_rel, stored_response_path, - knowledgebase, link_request, trace + knowledgebase, link_request, trace, + cpe ) VALUES ( $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, @@ -213,7 +226,8 @@ func (p *postgresDatabase) InsertBatch(ctx context.Context, results []runner.Res $48, $49, $50, $51, $52, $53, $54, $55, $56, $57, $58, $59, - $60, $61, $62 + $60, $61, $62, + $63 )`, tableName) stmt, err := tx.PrepareContext(ctx, query) @@ -235,6 +249,7 @@ func (p *postgresDatabase) InsertBatch(ctx context.Context, results []runner.Res kbJSON, _ := json.Marshal(r.KnowledgeBase) linkReqJSON, _ := json.Marshal(r.LinkRequest) traceJSON, _ := json.Marshal(r.Trace) + cpeJSON, _ := json.Marshal(r.CPE) _, err = stmt.ExecContext(ctx, r.Timestamp, r.URL, r.Input, r.Host, r.Port, r.Scheme, r.Path, r.Method, r.FinalURL, @@ -248,6 +263,7 @@ func (p *postgresDatabase) InsertBatch(ctx context.Context, results []runner.Res chainJSON, pq.Array(r.ChainStatusCodes), r.HeadlessBody, r.ScreenshotBytes, r.ScreenshotPath, r.ScreenshotPathRel, r.StoredResponsePath, kbJSON, linkReqJSON, traceJSON, + cpeJSON, ) if err != nil { return fmt.Errorf("failed to insert result: %w", err) diff --git a/runner/options.go b/runner/options.go index 50a08e773..4d1cc71c9 100644 --- a/runner/options.go +++ b/runner/options.go @@ -801,11 +801,10 @@ func (options *Options) ValidateOptions() error { var resolvers []string for _, resolver := range options.Resolvers { if fileutil.FileExists(resolver) { - chFile, err := fileutil.ReadFile(resolver) - if err != nil { - return errors.Wrapf(err, "Couldn't process resolver file \"%s\"", resolver) - } - for line := range chFile { + for line, err := range fileutil.Lines(resolver) { + if err != nil { + return errors.Wrapf(err, "Couldn't process resolver file \"%s\"", resolver) + } line = strings.TrimSpace(line) if line != "" && strings.Contains(line, ",") { for item := range strings.SplitSeq(line, ",") { diff --git a/runner/runner.go b/runner/runner.go index 0b7d4597a..af9db5661 100644 --- a/runner/runner.go +++ b/runner/runner.go @@ -743,11 +743,10 @@ func (r *Runner) streamInput() (chan string, error) { return } } else { - fchan, err := fileutil.ReadFile(r.options.InputFile) - if err != nil { - return - } - for item := range fchan { + for item, err := range fileutil.Lines(r.options.InputFile) { + if err != nil { + return + } if r.options.SkipDedupe || r.testAndSet(item) { if !trySend(item) { return @@ -761,11 +760,10 @@ func (r *Runner) streamInput() (chan string, error) { gologger.Fatal().Msgf("No input provided: %s", err) } for _, file := range files { - fchan, err := fileutil.ReadFile(file) - if err != nil { - return - } - for item := range fchan { + for item, err := range fileutil.Lines(file) { + if err != nil { + return + } if r.options.SkipDedupe || r.testAndSet(item) { if !trySend(item) { return @@ -775,11 +773,10 @@ func (r *Runner) streamInput() (chan string, error) { } } if fileutil.HasStdin() { - fchan, err := fileutil.ReadFileWithReader(os.Stdin) - if err != nil { - return - } - for item := range fchan { + for item, err := range fileutil.LinesReader(os.Stdin) { + if err != nil { + return + } if r.options.SkipDedupe || r.testAndSet(item) { if !trySend(item) { return