Land rapid7#7735, make assigning payloads fast again!

Author: wvu

lib/msf/core/exploit.rb

@@ -698,6 +698,12 @@ def target_arch
     (target and target.arch) ? target.arch : (arch == []) ? nil : arch
   end
 
+  def normalize_platform_arch
+    c_platform = (target && target.platform) ? target.platform : platform
+    c_arch     = (target && target.arch)     ? target.arch     : (arch == []) ? nil : arch
+    c_arch   ||= [ ARCH_X86 ]
+    return c_platform, c_arch
+  end
 
   #
   # Returns whether the requested payload is compatible with the module.
@@ -706,10 +712,23 @@ def target_arch
   # @return [TrueClass] Payload is compatible.
   # @return [FalseClass] Payload is not compatible.
   #
-  def is_payload_compatible?(payload_name)
-    payload_names = compatible_payloads.collect { |entry| entry[0] }
+  def is_payload_compatible?(name)
+    p = framework.payloads[name]
+
+    # Skip over payloads that are too big
+    return false if payload_space && p.cached_size && p.cached_size > payload_space
+
+    pi = p.new
 
-    payload_names.include?(payload_name)
+    # Are we compatible in terms of conventions and connections and
+    # what not?
+    return false if !compatible?(pi)
+
+    # If the payload is privileged but the exploit does not give
+    # privileged access, then fail it.
+    return false if !self.privileged && pi.privileged
+
+    return true
   end
 
   #
@@ -719,34 +738,11 @@ def is_payload_compatible?(payload_name)
   def compatible_payloads
     payloads = []
 
-
-    c_platform = (target and target.platform) ? target.platform : platform
-    c_arch     = (target and target.arch)     ? target.arch     : (arch == []) ? nil : arch
-    c_arch   ||= [ ARCH_X86 ]
+    c_platform, c_arch = normalize_platform_arch
 
     framework.payloads.each_module(
-      'Platform' => c_platform,
-      'Arch'     => c_arch ) { |name, mod|
-
-      # Skip over payloads that are too big
-      if ((payload_space) and
-          (framework.payloads.sizes[name]) and
-          (framework.payloads.sizes[name] > payload_space))
-        dlog("#{refname}: Skipping payload #{name} for being too large", 'core',
-          LEV_1)
-        next
-      end
-
-      # Are we compatible in terms of conventions and connections and
-      # what not?
-      next if (compatible?(framework.payloads.instance(name)) == false)
-
-      # If the payload is privileged but the exploit does not give
-      # privileged access, then fail it.
-      next if (self.privileged == false and framework.payloads.instance(name).privileged == true)
-
-      # This one be compatible!
-      payloads << [ name, mod ]
+      'Arch' => c_arch, 'Platform' => c_platform) { |name, mod|
+      payloads << [ name, mod ] if is_payload_compatible?(name)
     }
 
     return payloads;
@@ -758,12 +754,10 @@ def compatible_payloads
   def compatible_encoders
     encoders = []
 
-    c_platform = (target and target.platform) ? target.platform : platform
-    c_arch     = (target and target.arch)     ? target.arch     : (arch == []) ? nil : arch
+    c_platform, c_arch = normalize_platform_arch
 
     framework.encoders.each_module_ranked(
       'Arch' => c_arch, 'Platform' => c_platform) { |name, mod|
-
       encoders << [ name, mod ]
     }