From 442f36e4664fd0aaed648075c62cfc793ec14c48 Mon Sep 17 00:00:00 2001
From: William Vu <William_Vu@rapid7.com>
Date: Mon, 16 Dec 2019 18:35:52 -0600
Subject: [PATCH] Complete refactor of CheckModule

---
 lib/msf/core/exploit/check_module.rb | 60 +++++++++++-----------------
 1 file changed, 24 insertions(+), 36 deletions(-)

diff --git a/lib/msf/core/exploit/check_module.rb b/lib/msf/core/exploit/check_module.rb
index 1e96ebc41dc9..c88fe960bdcb 100644
--- a/lib/msf/core/exploit/check_module.rb
+++ b/lib/msf/core/exploit/check_module.rb
@@ -3,8 +3,6 @@
 #
 # This mixin implements an exploit's check method by invoking an aux module
 #
-# NOTE: The module's run_host/run method MUST return an Msf::Exploit::CheckCode
-#
 
 module Msf
 module Exploit::Remote::CheckModule
@@ -29,45 +27,35 @@ def check
       return CheckCode::Unsupported("Could not instantiate #{check_module}")
     end
 
-    # Bail if run_host/run isn't defined
-    if mod.respond_to?(:run_host)
-      meth = :run_host
-    elsif mod.respond_to?(:run)
-      meth = :run
-    else
-      return CheckCode::Unsupported("#{check_module} does not define a run_host/run method")
+    # Bail if it isn't aux
+    if mod.type != Msf::MODULE_AUX
+      return CheckCode::Unsupported("#{check_module} is not an auxiliary module")
     end
 
-    # Add the exploit's targeting options to the module's datastore
-    %w[RHOSTS RHOST RPORT].each do |opt|
-      next unless datastore[opt]
-
-      mod.datastore[opt] = datastore[opt].dup
+    # Bail if run isn't defined
+    unless mod.respond_to?(:run)
+      return CheckCode::Unsupported("#{check_module} does not define a run method")
     end
 
-    # Bail if module options don't validate
-    mod.options.validate(mod.datastore)
-
-    # Use the exploit's input and output as the module's
-    mod.user_input, mod.user_output = user_input, user_output
-
-    # Use the module's CheckCode
-    checkcode =
-      case meth
-      when :run_host
-        mod.run_host(rhost)
-      when :run
-        mod.run
-      end
-
-    # Bail if module doesn't return a CheckCode
-    unless checkcode.kind_of?(Exploit::CheckCode)
-      print_warning("#{check_module} does not return a CheckCode")
-      return Exploit::CheckCode::Unsupported
+    # Retrieve the module's return value
+    checkcode = mod.run_simple(
+      'LocalInput'  => user_input,
+      'LocalOutput' => user_output,
+      'Options'     => datastore.to_h.slice('RHOSTS', 'RHOST', 'RPORT')
+    )
+
+    # Ensure return value is a CheckCode
+    case checkcode
+    when Exploit::CheckCode
+      # Return the CheckCode
+      checkcode
+    when Hash
+      # XXX: Return scanner's last CheckCode
+      checkcode.values.last
+    else
+      # Bail if module doesn't return a CheckCode
+      Exploit::CheckCode::Unsupported("#{check_module} does not return a CheckCode")
     end
-
-    # Return the CheckCode
-    checkcode
   end
 
   def check_module