From a1c755161a6cbe5b879981d7856ff1f3b36d0061 Mon Sep 17 00:00:00 2001
From: HD Moore <hd_moore@rapid7.com>
Date: Sun, 29 Mar 2015 11:52:06 -0500
Subject: [PATCH] Add spec coverage for appender, fix injector

---
 .../lib/msf/core/exe/segment_appender_spec.rb | 82 +++++++++++++++++++
 .../lib/msf/core/exe/segment_injector_spec.rb |  6 --
 2 files changed, 82 insertions(+), 6 deletions(-)
 create mode 100644 spec/lib/msf/core/exe/segment_appender_spec.rb

diff --git a/spec/lib/msf/core/exe/segment_appender_spec.rb b/spec/lib/msf/core/exe/segment_appender_spec.rb
new file mode 100644
index 0000000000..5725eba822
--- /dev/null
+++ b/spec/lib/msf/core/exe/segment_appender_spec.rb
@@ -0,0 +1,82 @@
+require 'spec_helper'
+require 'msf/core/exe/segment_appender'
+
+describe Msf::Exe::SegmentAppender do
+
+  let(:opts) do
+    option_hash = {
+        :template => File.join(File.dirname(__FILE__), "..", "..", "..", "..", "..", "data", "templates", "template_x86_windows.exe"),
+        :payload  => "\xd9\xeb\x9b\xd9\x74\x24",
+        :arch     => :x86
+    }
+  end
+  subject(:injector) { Msf::Exe::SegmentInjector.new(opts) }
+
+  it { should respond_to :payload }
+  it { should respond_to :template }
+  it { should respond_to :arch }
+  it { should respond_to :processor }
+  it { should respond_to :buffer_register }
+
+  it 'should return the correct processor for the arch' do
+    injector.processor.class.should == Metasm::Ia32
+    injector.arch = :x64
+    injector.processor.class.should == Metasm::X86_64
+  end
+
+  context '#create_thread_stub' do
+    it 'should use edx as a default buffer register' do
+      injector.buffer_register.should == 'edx'
+    end
+
+    context 'when given a non-default buffer register' do
+      let(:opts) do
+        option_hash = {
+            :template => File.join(File.dirname(__FILE__), "..", "..", "..", "..", "..", "data", "templates", "template_x86_windows.exe"),
+            :payload  => "\xd9\xeb\x9b\xd9\x74\x24",
+            :arch     => :x86,
+            :buffer_register => 'eax'
+        }
+      end
+      it 'should use the correct buffer register' do
+        injector.buffer_register.should == 'eax'
+      end
+    end
+  end
+
+  describe '#generate_pe' do
+    it 'should return a string' do
+      injector.generate_pe.kind_of?(String).should == true
+    end
+
+    it 'should produce a valid PE exe' do
+      expect {Metasm::PE.decode(injector.generate_pe) }.to_not raise_exception
+    end
+
+    context 'the generated exe' do
+      let(:exe) { Metasm::PE.decode(injector.generate_pe) }
+      it 'should be the propper arch' do
+        exe.bitsize.should == 32
+      end
+
+      it 'should have 5 sections' do
+        exe.sections.count.should == 5
+      end
+
+      it 'should have all the right original section names' do
+        s_names = []
+        exe.sections.collect {|s| s_names << s.name}
+        s_names[0,4].should == [".text", ".rdata", ".data", ".rsrc"]
+      end
+
+      it 'should have the last section set to RWX' do
+        exe.sections.last.characteristics.should == ["CONTAINS_CODE", "MEM_EXECUTE", "MEM_READ", "MEM_WRITE"]
+      end
+
+      it 'should have an entrypoint that points to the last section' do
+        exe.optheader.entrypoint.should == exe.sections.last.virtaddr
+      end
+    end
+  end
+end
+
diff --git a/spec/lib/msf/core/exe/segment_injector_spec.rb b/spec/lib/msf/core/exe/segment_injector_spec.rb
index e5a4e9181a..3dd710cee5 100644
--- a/spec/lib/msf/core/exe/segment_injector_spec.rb
+++ b/spec/lib/msf/core/exe/segment_injector_spec.rb
@@ -24,12 +24,6 @@ describe Msf::Exe::SegmentInjector do
     injector.processor.class.should == Metasm::X86_64
   end
 
-  context '#payload_as_asm' do
-    it 'should return the payload as declare byte instructions' do
-      injector.payload_as_asm.should == "db 0xd9\ndb 0xeb\ndb 0x9b\ndb 0xd9\ndb 0x74\ndb 0x24\n"
-    end
-  end
-
   context '#create_thread_stub' do
     it 'should use edx as a default buffer register' do
       injector.buffer_register.should == 'edx'