Skip to content

Commit

Permalink
Add support for multiple output files
Browse files Browse the repository at this point in the history
  • Loading branch information
AI-Mozi committed Aug 27, 2024
1 parent 2338010 commit 7446b66
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 9 deletions.
21 changes: 13 additions & 8 deletions lib/ronin/recon/cli/commands/run.rb
Original file line number Diff line number Diff line change
Expand Up @@ -154,8 +154,7 @@ class Run < Command
usage: 'FILE'
},
desc: 'The output file to write results to' do |path|
options[:output] = path
options[:output_format] ||= OutputFormats.infer_from(path)
@outputs << [path, options[:output_format] || OutputFormats.infer_from(path)]
end

option :output_format, short: '-F',
Expand Down Expand Up @@ -229,6 +228,11 @@ class Run < Command
# @return [Array<Value>]
attr_reader :ignore

# The output files and formats
#
# @return [Array<(String, Class)>]
attr_reader :outputs

#
# Initializes the `ronin-recon run` command.
#
Expand All @@ -246,7 +250,8 @@ def initialize(**kwargs)
@worker_params = Hash.new { |hash,key| hash[key] = {} }
@worker_concurrency = {}

@ignore = []
@ignore = []
@outputs = []
end

#
Expand All @@ -261,9 +266,9 @@ def run(*values)

values = values.map { |value| parse_value(value) }

output_file = if options[:output] && options[:output_format]
options[:output_format].open(options[:output])
end
output_files = outputs.filter_map do |output, output_format|
output_format.open(output) if output
end

if options[:import]
require 'ronin/db'
Expand All @@ -280,7 +285,7 @@ def run(*values)
print_value(value,parent)
end

if output_file
output_files.each do |output_file|
engine.on(:value) do |value|
output_file << value
end
Expand Down Expand Up @@ -308,7 +313,7 @@ def run(*values)
end
end
ensure
output_file.close if options[:output]
output_files&.each(&:close)
end
end

Expand Down
4 changes: 3 additions & 1 deletion spec/cli/commands/run_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,9 @@
end

it "must set the :output_format option using the path's file extension" do
expect(subject.options[:output_format]).to be(Ronin::Core::OutputFormats::JSON)
expect(subject.outputs.size).to eq(1)
expect(subject.outputs[0][0]).to eq(path)
expect(subject.outputs[0][1]).to be(Ronin::Core::OutputFormats::JSON)
end

context "but the '--output-format' has already been specified" do
Expand Down

0 comments on commit 7446b66

Please sign in to comment.