generating raw payload bits now

added raw payload generation, arch selection,
and specs for everything thus far

lib/msf/core/payload_generator.rb

@@ -3,6 +3,9 @@ module Msf
   class IncompatiblePlatform < StandardError
   end
 
+  class IncompatibleArch < StandardError
+  end
+
   class PayloadGenerator
 
 
@@ -19,10 +22,11 @@ module Msf
     attr_accessor :payload
     attr_accessor :platform
     attr_accessor :space
+    attr_accessor :stdin
     attr_accessor :template
 
 
-    # @param [Hash] opts The options hash
+    # @param opts [Hash] The options hash
     # @option opts [String] :payload The refname of the payload to generate
     # @option opts [String] :format The format you want the payload returned in
     # @option opts [String] :encoder The encoder you want applied to the payload
@@ -38,7 +42,7 @@ module Msf
     # @option opts [Hash] :datastore The datastore to apply to the payload module
     # @option opts [Msf::Framework] :framework The framework instance to use for generation
     def initialize(opts={})
-      @add_code   = opts.fetch(:add_code, false)
+      @add_code   = opts.fetch(:add_code, '')
       @arch       = opts.fetch(:arch, '')
       @badchars   = opts.fetch(:badchars, '')
       @datastore  = opts.fetch(:datastore, {})
@@ -50,6 +54,7 @@ module Msf
       @payload    = opts.fetch(:payload, '')
       @platform   = opts.fetch(:platform, '')
       @space      = opts.fetch(:size, 1073741824)
+      @stdin      = opts.fetch(:stdin, nil)
       @template   = opts.fetch(:template, '')
 
       @framework  = opts.fetch(:framework)
@@ -58,15 +63,53 @@ module Msf
       raise ArgumentError, "Invalid Format Selected" unless format_is_valid?
     end
 
-
+    # @raise [Msf::IncompatiblePlatform] if no platform was selected for a stdin payload
+    # @raise [Msf::IncompatibleArch] if no arch was selected for a stdin payload
+    # @raise [Msf::IncompatiblePlatform] if the platform is incompatible with the payload
+    # @raise [Msf::IncompatibleArch] if the arch is incompatible with the payload
+    # @return [String] the raw bytes of the payload to be generated
     def generate_raw_payload
+      if payload == 'stdin'
+        if arch.blank?
+          raise IncompatibleArch, "You must select an arch for a custom payload"
+        elsif platform.blank?
+          raise IncompatiblePlatform, "You must select a platform for a custom payload"
+        end
+        stdin
+      else
         payload_module = framework.payloads.create(payload)
 
         chosen_platform = choose_platform(payload_module)
         if chosen_platform.platforms.empty?
-        raise IncompatiblePlatform, "The selected platform is Incompatible with the Payload"
+          raise IncompatiblePlatform, "The selected platform is incompatible with the payload"
         end
 
+        chosen_arch = choose_arch(payload_module)
+        unless chosen_arch
+          raise IncompatibleArch, "The selected arch is incompatible with the payload"
+        end
+
+        payload_module.generate_simple(
+            'Format'   => 'raw',
+            'Options'  => datastore,
+            'Encoder'  => nil
+        )
+      end
+    end
+
+    # @param shellcode [String] The shellcode to add to
+    # @return [String] the combined shellcode which executes the added code in a seperate thread
+    def add_shellcode(shellcode)
+      if add_code.present? and platform_list.platforms.include? Msf::Module::Platform::Windows and arch == "x86"
+        shellcode_file = File.open(add_code)
+        shellcode_file.binmode
+        added_code = shellcode_file.read
+        shellcode_file.close
+        shellcode = ::Msf::Util::EXE.win32_rwx_exec_thread(shellcode,0,'end')
+        shellcode << added_code
+      else
+        shellcode.dup
+      end
     end
 
     # @return [Msf::Module::PlatformList] It will be empty if no valid platforms found
@@ -83,22 +126,40 @@ module Msf
       list
     end
 
+    # @param mod [Msf::Module] The module class to choose a platform for
+    # @return [Msf::Module::PlatformList] The selected platform list
     def choose_platform(mod)
       chosen_platform = platform_list
       if chosen_platform.platforms.empty?
         chosen_platform = mod.platform
+        @platform = mod.platform.platforms.first.to_s.split("::").last
       elsif (chosen_platform & mod.platform).empty?
         chosen_platform = Msf::Module::PlatformList.new
       end
       chosen_platform
     end
 
+    # @param mod [Msf::Module] The module class to choose an arch for
+    # @return [String] String form of the Arch if a valid arch found
+    # @return [Nil] if no valid arch found
+    def choose_arch(mod)
+      if arch.blank?
+        @arch = mod.arch.first
+        return mod.arch.first
+      elsif mod.arch.include? arch
+        return arch
+      else
+        return nil
+      end
+    end
+
+
     private
 
     # @return [True] if the payload is a valid Metasploit Payload
     # @return [False] if the payload is not a valid Metasploit Payload
     def payload_is_valid?
-      framework.payloads.keys.include? payload
+      (framework.payloads.keys + ['stdin']).include? payload
     end
 
     # @return [True] if the format is valid

spec/lib/msf/core/payload_generator_spec.rb

@@ -24,6 +24,7 @@ describe Msf::PayloadGenerator do
   let(:payload) { "windows/meterpreter/reverse_tcp"}
   let(:platform) { "Windows" }
   let(:space) { 1073741824 }
+  let(:stdin) { nil }
   let(:template) { File.join(Msf::Config.data_directory, "templates", "template_x86_windows.exe") }
   let(:generator_opts) {
     {
@@ -40,9 +41,11 @@ describe Msf::PayloadGenerator do
         payload: payload,
         platform: platform,
         space: space,
+        stdin: stdin,
         template: template
     }
   }
+  let(:payload_module) { framework.payloads.create(payload)}
 
   subject(:payload_generator) { described_class.new(generator_opts) }
 
@@ -59,6 +62,7 @@ describe Msf::PayloadGenerator do
   it { should respond_to :payload }
   it { should respond_to :platform }
   it { should respond_to :space }
+  it { should respond_to :stdin }
   it { should respond_to :template }
 
   context 'when creating a new generator' do
@@ -79,6 +83,7 @@ describe Msf::PayloadGenerator do
             payload: payload,
             platform: platform,
             space: space,
+            stdin: stdin,
             template: template
         }
       }
@@ -98,6 +103,12 @@ describe Msf::PayloadGenerator do
       it { should raise_error(ArgumentError, "Invalid Payload Selected") }
     end
 
+    context 'when given a payload through stdin' do
+      let(:payload) { "stdin" }
+
+      it { should_not raise_error }
+    end
+
     context 'when not given a format' do
       let(:format) { nil }
 
@@ -123,9 +134,6 @@ describe Msf::PayloadGenerator do
     end
   end
 
-  context 'checking platforms' do
-    let(:payload_module) { framework.payloads.create(payload)}
-
   context 'when not given a platform' do
     let(:platform) { '' }
 
@@ -139,6 +147,12 @@ describe Msf::PayloadGenerator do
       it 'chooses the platform list for the module' do
         expect(payload_generator.choose_platform(payload_module).platforms).to eq [Msf::Module::Platform::Windows]
       end
+
+      it 'sets the platform attr to the first platform of the module' do
+        my_generator = payload_generator
+        my_generator.choose_platform(payload_module)
+        expect(my_generator.platform).to eq "Windows"
+      end
     end
 
   end
@@ -185,6 +199,84 @@ describe Msf::PayloadGenerator do
 
   end
 
+  context '#choose_arch' do
+    context 'when no arch is selected' do
+      let(:arch) { '' }
+
+      it 'returns the first arch of the module' do
+        expect(payload_generator.choose_arch(payload_module)).to eq "x86"
+      end
+
+      it 'sets the arch to match the module' do
+        my_generator = payload_generator
+        my_generator.choose_arch(payload_module)
+        expect(my_generator.arch).to eq "x86"
+      end
+    end
+
+    context 'when the arch matches the module' do
+      it 'returns the selected arch' do
+        expect(payload_generator.choose_arch(payload_module)).to eq arch
+      end
+    end
+
+    context 'when the arch does not match the module' do
+      let(:arch) { "mipsle" }
+
+      it "returns nil" do
+        expect(payload_generator.choose_arch(payload_module)).to be_nil
+      end
+    end
+  end
+
+  context '#generate_raw_payload' do
+
+    context 'when passing a payload through stdin' do
+      let(:stdin) { "\x90\x90\x90"}
+      let(:payload) { "stdin" }
+
+      context 'when no arch has been selected' do
+        let(:arch) { '' }
+
+        it 'raises an IncompatibleArch error' do
+          expect{payload_generator.generate_raw_payload}.to raise_error(Msf::IncompatibleArch, "You must select an arch for a custom payload")
+        end
+      end
+
+      context 'when no platform has been selected' do
+        let(:platform) { '' }
+
+        it 'raises an IncompatiblePlatform error' do
+          expect{payload_generator.generate_raw_payload}.to raise_error(Msf::IncompatiblePlatform, "You must select a platform for a custom payload")
+        end
+      end
+
+      it 'returns the payload from stdin' do
+        expect(payload_generator.generate_raw_payload).to eq stdin
+      end
+    end
+
+    context 'when selecting a metasploit payload' do
+      context 'when the platform is incompatible with the payload' do
+        let(:platform) { "linux" }
+
+        it 'raises an IncompatiblePlatform error' do
+          expect{payload_generator.generate_raw_payload}.to raise_error(Msf::IncompatiblePlatform, "The selected platform is incompatible with the payload")
+        end
+      end
+
+      context 'when the arch is incompatible with the payload' do
+        let(:arch) { "mipsle" }
+
+        it 'raises an IncompatibleArch error' do
+          expect{payload_generator.generate_raw_payload}.to raise_error(Msf::IncompatibleArch, "The selected arch is incompatible with the payload")
+        end
+      end
+
+      it 'returns the raw bytes of the payload' do
+        expect(payload_generator.generate_raw_payload).to be_present
+      end
+    end
   end