Land #12509, add check result to RPC API

lib/msf/base/simple/auxiliary.rb

@@ -68,21 +68,25 @@ module Auxiliary
       mod.init_ui(nil, nil)
     end
 
-    ctx = [ mod ]
+    run_uuid = Rex::Text.rand_text_alphanumeric(24)
+    mod.framework.ready << run_uuid
+    ctx = [mod, run_uuid]
     if(mod.passive? or opts['RunAsJob'])
       mod.job_id = mod.framework.jobs.start_bg_job(
         "Auxiliary: #{mod.refname}",
         ctx,
-        Proc.new { |ctx_| self.job_run_proc(ctx_) },
+        Proc.new { |ctx_| self.job_run_proc(ctx_, &:run) },
         Proc.new { |ctx_| self.job_cleanup_proc(ctx_) }
       )
       # Propagate this back to the caller for console mgmt
       omod.job_id = mod.job_id
+      return [run_uuid, mod.job_id]
     else
-      self.job_run_proc(ctx)
+      result = self.job_run_proc(ctx, &:run)
       self.job_cleanup_proc(ctx)
-    end
 
+      return result
+    end
   end
 
   #
@@ -105,6 +109,9 @@ module Auxiliary
   # 	The local output through which data can be displayed.
   #
   def self.check_simple(mod, opts)
+    Msf::Simple::Framework.simplify_module(mod, false)
+
+    mod._import_extra_options(opts)
     if opts['LocalInput']
       mod.init_ui(opts['LocalInput'], opts['LocalOutput'])
     end
@@ -113,10 +120,33 @@ module Auxiliary
     # be normalized
     mod.validate
 
-    mod.setup
 
-    # Run check if it exists
-    mod.respond_to?(:check) ? mod.check : Msf::Exploit::CheckCode::Unsupported
+    run_uuid = Rex::Text.rand_text_alphanumeric(24)
+    mod.framework.ready << run_uuid
+    ctx = [mod, run_uuid]
+
+    if opts['RunAsJob']
+      mod.job_id = mod.framework.jobs.start_bg_job(
+        "Auxiliary: #{mod.refname} check",
+        ctx,
+        Proc.new do |ctx_|
+          self.job_run_proc(ctx_) do |m|
+            m.respond_to?(:check) ? m.check : Msf::Exploit::CheckCode::Unsupported
+          end
+        end,
+        Proc.new { |ctx_| self.job_cleanup_proc(ctx_) }
+      )
+
+      [run_uuid, mod.job_id]
+    else
+      # Run check if it exists
+      result = self.job_run_proc(ctx) do |m|
+        m.respond_to?(:check) ? m.check : Msf::Exploit::CheckCode::Unsupported
+      end
+      self.job_cleanup_proc(ctx)
+
+      result
+    end
   end
 
   #
@@ -132,12 +162,23 @@ protected
   #
   # Job run proc, sets up the module and kicks it off.
   #
-  def self.job_run_proc(ctx)
+  def self.job_run_proc(ctx, &block)
     mod = ctx[0]
+    run_uuid = ctx[1]
     begin
       mod.setup
       mod.framework.events.on_module_run(mod)
-      mod.run
+      begin
+        mod.framework.running << run_uuid
+        mod.framework.ready.delete run_uuid
+        result = block.call(mod)
+        mod.framework.results[run_uuid] = {result: result}
+      rescue Exception => e
+        mod.framework.results[run_uuid] = {error: e.to_s}
+        raise
+      ensure
+        mod.framework.running.delete run_uuid
+      end
     rescue Msf::Auxiliary::Complete
       mod.cleanup
       return
@@ -178,8 +219,8 @@ protected
 
       mod.cleanup
 
-      return
     end
+    return result
   end
 
   #

lib/msf/base/simple/exploit.rb

@@ -178,6 +178,9 @@ module Exploit
   # 	The local output through which data can be displayed.
   #
   def self.check_simple(mod, opts)
+    Msf::Simple::Framework.simplify_module(mod, false)
+    mod._import_extra_options(opts)
+
     if opts['LocalInput']
       mod.init_ui(opts['LocalInput'], opts['LocalOutput'])
     end
@@ -186,10 +189,21 @@ module Exploit
     # be normalized
     mod.validate
 
-    mod.setup
+    run_uuid = Rex::Text.rand_text_alphanumeric(24)
+    mod.framework.ready << run_uuid
+    ctx = [mod, run_uuid]
 
-    # Run check if it exists
-    mod.respond_to?(:check) ? mod.check : Msf::Exploit::CheckCode::Unsupported
+    if opts['RunAsJob']
+      mod.job_id = mod.framework.jobs.start_bg_job(
+        "Auxiliary: #{mod.refname} check",
+        ctx,
+        Proc.new { |ctx_| self.job_check_proc(ctx_) },
+        Proc.new { |ctx_| nil }
+      )
+      [run_uuid, mod.job_id]
+    else
+      self.job_check_proc(ctx)
+    end
   end
 
   #
@@ -199,6 +213,26 @@ module Exploit
     Msf::Simple::Exploit.check_simple(self, opts)
   end
 
+  protected
+
+  def self.job_check_proc(ctx)
+    mod = ctx[0]
+    run_uuid = ctx[1]
+    begin
+      mod.setup
+      mod.framework.running << run_uuid
+      mod.framework.ready.delete run_uuid
+      result = mod.respond_to?(:check) ? mod.check : Msf::Exploit::CheckCode::Unsupported
+      mod.framework.results[run_uuid] = {result: result}
+    rescue => e
+      mod.framework.results[run_uuid] = {error: e.to_s}
+      mod.handle_exception e
+    ensure
+      mod.framework.running.delete run_uuid
+    end
+
+    return result
+  end
 end
 
 end

lib/msf/base/simple/framework.rb

@@ -150,6 +150,9 @@ module Framework
   #
   def init_simplified
     self.stats = Statistics.new(self)
+    self.ready = Set.new
+    self.running = Set.new
+    self.results = Hash.new
   end
 
   #
@@ -180,12 +183,31 @@ module Framework
   # Thread of the running rebuild operation
   #
   attr_reader :cache_thread
+
+  #
+  # {Set<String>} of module run/check UUIDs waiting to be kicked off
+  #
+  attr_reader :ready
+
+  #
+  # {Hash<String,Hash>} of module run/check results, by UUID. Successful runs
+  # look like `{result: check_code}` and errors like `{error: message}`.
+  #
+  attr_reader :results
+
+  #
+  # {Set<String>} of module run/check UUIDs currently in progress
+  #
+  attr_reader :running
   attr_writer :cache_initialized # :nodoc:
   attr_writer :cache_thread # :nodoc:
 
 
 protected
 
+  attr_writer :ready # :nodoc:
+  attr_writer :results # :nodoc:
+  attr_writer :running # :nodoc:
   attr_writer :stats # :nodoc:
 
 end

lib/msf/core/auxiliary/scanner.rb

@@ -60,6 +60,9 @@ def run
   @tl = []
   @scan_errors = []
 
+  res = Queue.new
+  results = Hash.new
+
   #
   # Sanity check threading given different conditions
   #
@@ -108,7 +111,7 @@ def run
           nmod.datastore['RHOST'] = targ
 
           begin
-            nmod.run_host(targ)
+            res << {tip => nmod.run_host(targ)}
           rescue ::Rex::BindFailed
             if datastore['CHOST']
               @scan_errors << "The source IP (CHOST) value of #{datastore['CHOST']} was not usable"
@@ -125,6 +128,11 @@ def run
         end
       end
 
+      # Do as much of this work as possible while other threads are running
+      while !res.empty?
+        results.merge! res.pop
+      end
+
       # Stop scanning if we hit a fatal error
       break if has_fatal_errors?
 
@@ -148,7 +156,7 @@ def run
     end
 
     scanner_handle_fatal_errors
-    return
+    return results
   end
 
   if (self.respond_to?('run_batch'))

lib/msf/core/rpc/v10/rpc_module.rb

@@ -483,6 +483,57 @@ class RPC_Module < RPC_Base
 
   end
 
+  # Runs the check method of a module.
+  #
+  # @param [String] mtype Module type. Supported types include (case-sensitive):
+  #                       * exploit
+  #                       * auxiliary
+  # @param [String] mname Module name. For example: 'windows/smb/ms08_067_netapi'.
+  # @param [Hash] opts Options for the module (such as datastore options).
+  # @raise [Msf::RPC::Exception] Module not found (either wrong type or name).
+  # @return
+  def rpc_check(mtype, mname, opts)
+    mod = _find_module(mtype,mname)
+    case mtype
+    when 'exploit'
+      _check_exploit(mod, opts)
+    when 'auxiliary'
+      _run_auxiliary(mod, opts)
+    else
+      error(500, "Invalid Module Type: #{mtype}")
+    end
+  end
+
+  # TODO: expand these to take a list of UUIDs or stream with event data if
+  # required for performance
+  def rpc_results(uuid)
+    if r = self.framework.results[uuid]
+      if r[:error]
+        {"status" => "errored", "error" => r[:error]}
+      else
+        if r[:result].length == 1
+          # A hash of one IP => result
+          # TODO: make hashes of IP => result the normal case
+          {"status" => "completed", "result" => r[:result].values.first}
+        else
+          # Either singular check code or multiple hosts
+          # TODO: combine underlying code so that nothing returns a bare CheckCode anymore
+          {"status" => "completed", "result" => r[:result]}
+        end
+      end
+    elsif self.framework.running.include? uuid
+      {"status" => "running"}
+    elsif self.framework.ready.include? uuid
+      {"status" => "ready"}
+    else
+      error(404, "Results not found for module instance #{uuid}")
+    end
+  end
+
+  def rpc_ack(uuid)
+    {"success" => !!self.framework.results.delete(uuid)}
+  end
+
   # Returns a list of executable format names.
   #
   # @return [Array<String>] A list of executable format names, for example: ["exe"]
@@ -673,14 +724,37 @@ private
   end
 
   def _run_auxiliary(mod, opts)
-    Msf::Simple::Auxiliary.run_simple(mod, {
+    uuid, job = Msf::Simple::Auxiliary.run_simple(mod, {
       'Action'   => opts['ACTION'],
       'RunAsJob' => true,
       'Options'  => opts
     })
     {
-      "job_id" => mod.job_id,
-      "uuid" => mod.uuid
+      "job_id" => job,
+      "uuid" => uuid
+    }
+  end
+
+  def _check_exploit(mod, opts)
+    uuid, job = Msf::Simple::Exploit.check_simple(mod, {
+        'RunAsJob' => true,
+        'Options'  => opts
+    })
+    {
+      "job_id" => job,
+      "uuid" => uuid
+    }
+  end
+
+  def _check_auxiliary(mod, opts)
+    uuid, job = Msf::Simple::Auxiliary.check_simple(mod, {
+        'Action'   => opts['ACTION'],
+        'RunAsJob' => true,
+        'Options'  => opts
+    })
+    {
+      "job_id" => job,
+      "uuid" => uuid
     }
   end
 

spec/lib/msf/base/simple/framework_spec.rb

@@ -8,4 +8,52 @@ RSpec.describe Msf::Simple::Framework do
   end
 
   it_should_behave_like 'Msf::Simple::Framework::ModulePaths'
+
+  describe "#ready" do
+    let(:run_uuid) { Rex::Text.rand_text_alphanumeric 24 }
+    it "should start out empty" do
+      expect(subject.ready).to_not include, run_uuid
+    end
+    it "should remember things that are ready to run" do
+      subject.ready << run_uuid
+      expect(subject.ready).to include, run_uuid
+    end
+    it "should forget things that are running" do
+      subject.ready << run_uuid
+      subject.ready.delete run_uuid
+      expect(subject.ready).to_not include, run_uuid
+    end
+  end
+
+  describe "#running" do
+    let(:run_uuid) { Rex::Text.rand_text_alphanumeric 24 }
+    it "should start out empty" do
+      expect(subject.running).to_not include, run_uuid
+    end
+    it "should remember things that are running" do
+      subject.running << run_uuid
+      expect(subject.running).to include, run_uuid
+    end
+    it "should forget things that are done" do
+      subject.running << run_uuid
+      subject.running.delete run_uuid
+      expect(subject.running).to_not include, run_uuid
+    end
+  end
+
+  describe "#results" do
+    let(:run_uuid) { Rex::Text.rand_text_alphanumeric 24 }
+    it "should start out empty" do
+      expect(subject.results.keys).to_not include, run_uuid
+    end
+    it "should remember results" do
+      subject.results[run_uuid] = {}
+      expect(subject.results.keys).to include, run_uuid
+    end
+    it "should forget things that have been acknowleged" do
+      subject.results[run_uuid] = {}
+      subject.results.delete run_uuid
+      expect(subject.results.keys).to_not include, run_uuid
+    end
+  end
 end