Land rapid7#5116, better handle platform and arch in msfvenom

Author: Brent Cook

lib/msf/core/payload_generator.rb

@@ -144,8 +144,10 @@ def choose_arch(mod)
       if arch.blank?
         @arch = mod.arch.first
         cli_print "No Arch selected, selecting Arch: #{arch} from the payload"
+        datastore['ARCH'] = arch if mod.kind_of?(Msf::Payload::Generic)
         return mod.arch.first
       elsif mod.arch.include? arch
+        datastore['ARCH'] = arch if mod.kind_of?(Msf::Payload::Generic)
         return arch
       else
         return nil
@@ -157,14 +159,28 @@ def choose_arch(mod)
     # @param mod [Msf::Payload] The module class to choose a platform for
     # @return [Msf::Module::PlatformList] The selected platform list
     def choose_platform(mod)
+      # By default, platform_list will at least return Msf::Module::Platform
+      # if there is absolutely no pre-configured platform info at all
       chosen_platform = platform_list
+
       if chosen_platform.platforms.empty?
         chosen_platform = mod.platform
         cli_print "No platform was selected, choosing #{chosen_platform.platforms.first} from the payload"
         @platform = mod.platform.platforms.first.to_s.split("::").last
       elsif (chosen_platform & mod.platform).empty?
         chosen_platform = Msf::Module::PlatformList.new
       end
+
+      begin
+        platform_object = Msf::Module::Platform.find_platform(platform)
+      rescue ArgumentError
+        platform_object = nil
+      end
+
+      if mod.kind_of?(Msf::Payload::Generic) && mod.send(:module_info)['Platform'].empty? && platform_object
+        datastore['PLATFORM'] = platform
+      end
+
       chosen_platform
     end