From c44c0ebf48cf78dc73fd6681a5e22264bc6c6fb9 Mon Sep 17 00:00:00 2001
From: James Lee <egypt@metasploit.com>
Date: Fri, 2 Mar 2012 17:55:11 -0700
Subject: [PATCH] Refactor ModuleTest and add a few more tests

This makes running tests from a post module as easy as creating methods
that start with +test_+ and running +it()+ blocks inside them.
---
 test/lib/module_test.rb               | 37 ++++++++++-
 test/modules/post/test/meterpreter.rb | 96 ++++++++++++++++++++++-----
 2 files changed, 112 insertions(+), 21 deletions(-)

diff --git a/test/lib/module_test.rb b/test/lib/module_test.rb
index 5d95d9a945..eb60b90467 100644
--- a/test/lib/module_test.rb
+++ b/test/lib/module_test.rb
@@ -1,16 +1,33 @@
 
 
+module Msf
 
-module Msf::ModuleTest
-	attr_accessor :error
+module ModuleTest
+	attr_accessor :tests
+	attr_accessor :failures
+
+	def initialize(info={})
+		@tests = 0
+		@failures = 0
+		super
+	end
+
+	def run_all_tests
+		tests = self.methods.select { |m| m.to_s =~ /^test_/ }
+		tests.each { |test_method|
+			self.send(test_method)
+		}
+
+	end
 
 	def it(msg="", &block)
+		@tests  += 1
 		begin
 			result = block.call
 			unless result
 				print_error("FAILED: #{msg}")
 				print_error("FAILED: #{error}") if error
-				@error = nil
+				@failures  += 1
 				return
 			end
 		rescue ::Exception => e
@@ -23,5 +40,19 @@ module Msf::ModuleTest
 	end
 end
 
+module ModuleTest::PostTest
+	include ModuleTest
+	def run
+		print_status("Running against session #{datastore["SESSION"]}")
+		print_status("Session type is #{session.type}")
+		print_status("Session platform is #{session.platform}")
 
+		@tests = 0; @failures = 0
+		run_all_tests
 
+		vprint_status("Testing complete.")
+		print_status("Passed: #{@tests - @failures}; Failed: #{@failures}")
+	end
+end
+
+end
diff --git a/test/modules/post/test/meterpreter.rb b/test/modules/post/test/meterpreter.rb
index de1bf2fb4b..a0a2438ab1 100644
--- a/test/modules/post/test/meterpreter.rb
+++ b/test/modules/post/test/meterpreter.rb
@@ -6,11 +6,9 @@ $:.push "test/lib" unless $:.include? "test/lib"
 #require 'module_test'
 load 'test/lib/module_test.rb'
 
-load 'lib/rex/post/meterpreter/extensions/stdapi/fs/dir.rb'
-
 class Metasploit3 < Msf::Post
 
-	include Msf::ModuleTest
+	include Msf::ModuleTest::PostTest
 
 	def initialize(info={})
 		super( update_info( info,
@@ -18,25 +16,16 @@ class Metasploit3 < Msf::Post
 				'Description'   => %q{ This module will test meterpreter API methods },
 				'License'       => MSF_LICENSE,
 				'Author'        => [ 'egypt'],
-				'Version'       => '$Revision: 11663 $',
-				'Platform'      => [ 'windows', 'linux' ],
+				'Version'       => '$Revision$',
+				'Platform'      => [ 'windows', 'linux', 'java' ],
 				'SessionTypes'  => [ 'meterpreter' ]
 			))
 
 	end
 
-	def run
-		print_status("Running against session #{datastore["SESSION"]}")
-		print_status("Session type is #{session.type}")
-
-		test_sys_config
-		test_net_config
-		test_fs
-
-		print_status("Testing complete.")
-	end
-
 	def test_sys_config
+		vprint_status("Starting system config tests")
+
 		it "should return a user id" do
 			uid = session.sys.config.getuid
 			true
@@ -49,14 +38,49 @@ class Metasploit3 < Msf::Post
 	end
 
 	def test_net_config
+		vprint_status("Starting networking tests")
+
 		it "should return network interfaces" do
 			ifaces = session.net.config.get_interfaces
+			res = !!(ifaces and ifaces.length > 0)
 
-			ifaces and ifaces.length > 0
+			res
 		end
+		it "should have an interface that matches session_host" do
+			ifaces = session.net.config.get_interfaces
+			res = !!(ifaces and ifaces.length > 0)
+
+			p session.session_host
+			res &&= !! ifaces.find { |iface|
+				iface.ip == session.session_host || iface.ip6 == session.session_host
+			}
+
+			res
+		end
+
+		it "should return network routes" do
+			routes = session.net.config.get_routes
+
+			routes[0] and routes[0].length > 0
+		end
+
 	end
 
 	def test_fs
+		vprint_status("Starting filesystem tests")
+
+		it "should return the proper directory separator" do
+			sysinfo = session.sys.config.sysinfo
+			if sysinfo["OS"] =~ /windows/i
+				sep = session.fs.file.separator
+				res = (sep == "\\")
+			else
+				sep = session.fs.file.separator
+				res = (sep == "/")
+			end
+
+			res
+		end
 
 		it "should return the current working directory" do
 			wd = session.fs.dir.pwd
@@ -119,7 +143,6 @@ class Metasploit3 < Msf::Post
 			fd = session.fs.file.new("meterpreter-test", "rb")
 			contents = fd.read
 			vprint_status("Wrote #{contents}")
-			p fd
 			res &&= (contents == "test")
 			fd.close
 
@@ -129,8 +152,45 @@ class Metasploit3 < Msf::Post
 			res
 		end
 
+		it "should upload a file" do
+			res = true
+			remote = "HACKING.remote.txt"
+			local  = "HACKING"
+			vprint_status("uploading")
+			session.fs.file.upload_file(remote, local)
+			vprint_status("done")
+			res &&= session.fs.dir.entries.include?(remote)
+			vprint_status("remote file exists? #{res.inspect}")
+
+			if res
+				session.fs.file.download(remote, remote)
+				res &&= ::File.file? remote
+				::File.unlink remote
+			end
+
+			res
+		end
+
 	end
 
+	def test_sniffer
+		begin
+			session.core.use "sniffer"
+		rescue
+			# Not all meterpreters have a sniffer extension, don't count it
+			# against them.
+			return
+		end
+
+		it "should list interfaces" do
+			session.sniffer.interfaces.kind_of? Array
+		end
+
+		# XXX: how do we test this more thoroughly in a generic way?
+	end
+
+protected
+
 	def create_directory(name)
 		res = true