From 7446b66b8f69d8266596579c3a043d6106a3aee0 Mon Sep 17 00:00:00 2001
From: AI-Mozi <bieniek.milosz0@gmail.com>
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..c5d4d34 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<Value>]
           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) if 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