From b63b58346b4dbcf5a41a64800913c8b06857d777 Mon Sep 17 00:00:00 2001 From: AI-Mozi Date: Sun, 18 Aug 2024 16:44:11 +0200 Subject: [PATCH] Add support for multiple output files --- lib/ronin/recon/cli/commands/run.rb | 21 +++++++++++++-------- spec/cli/commands/run_spec.rb | 4 +++- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/lib/ronin/recon/cli/commands/run.rb b/lib/ronin/recon/cli/commands/run.rb index 21dc32c..28d581b 100644 --- a/lib/ronin/recon/cli/commands/run.rb +++ b/lib/ronin/recon/cli/commands/run.rb @@ -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', @@ -229,6 +228,11 @@ class Run < Command # @return [Array] attr_reader :ignore + # The output files and formats + # + # @return [Array<(String, Class)>] + attr_reader :outputs + # # Initializes the `ronin-recon run` command. # @@ -246,7 +250,8 @@ def initialize(**kwargs) @worker_params = Hash.new { |hash,key| hash[key] = {} } @worker_concurrency = {} - @ignore = [] + @ignore = [] + @outputs = [] end # @@ -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) + end if options[:import] require 'ronin/db' @@ -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 @@ -308,7 +313,7 @@ def run(*values) end end ensure - output_file.close if options[:output] + output_files&.each(&:close) end end diff --git a/spec/cli/commands/run_spec.rb b/spec/cli/commands/run_spec.rb index 1f3a735..4febdc8 100644 --- a/spec/cli/commands/run_spec.rb +++ b/spec/cli/commands/run_spec.rb @@ -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