From 3c40aba395bae9b40166fb5fca0fe00f728c8d63 Mon Sep 17 00:00:00 2001
From: Jonathan Cran <jcran@rapid7.com>
Date: Sat, 17 Sep 2011 03:00:57 +0000
Subject: [PATCH] updated to add improved esxi support

git-svn-id: file:///home/svn/framework3/trunk@13744 4d416f70-5f16-0410-b530-b9f4589650da
---
 lib/lab/driver/fog_driver.rb             |  2 -
 lib/lab/driver/remote_esx_driver.rb      | 62 ++++++++++++++++++++----
 lib/lab/driver/vm_driver.rb              | 23 ++++-----
 lib/lab/modifier/meterpreter_modifier.rb |  7 ++-
 lib/lab/modifier/test_modifier.rb        |  8 ---
 lib/lab/modifiers.rb                     |  1 +
 lib/lab/vm.rb                            | 20 +++++---
 7 files changed, 81 insertions(+), 42 deletions(-)

diff --git a/lib/lab/driver/fog_driver.rb b/lib/lab/driver/fog_driver.rb
index dde2c210ca..9a4a61e6ba 100644
--- a/lib/lab/driver/fog_driver.rb
+++ b/lib/lab/driver/fog_driver.rb
@@ -13,8 +13,6 @@ class FogDriver < VmDriver
 		super(config)
 		@fog_config = fog_config
 
-		puts "Fog Config: #{fog_config}"
-
 		# Soft dependency
 		begin
 			require 'fog'
diff --git a/lib/lab/driver/remote_esx_driver.rb b/lib/lab/driver/remote_esx_driver.rb
index 5cb7d37d94..66d4f18883 100644
--- a/lib/lab/driver/remote_esx_driver.rb
+++ b/lib/lab/driver/remote_esx_driver.rb
@@ -13,12 +13,11 @@ module Drivers
 class RemoteEsxDriver < VmDriver
 
 	def initialize(config)
-
 		unless config['user'] then raise ArgumentError, "Must provide a username" end
 		unless config['host'] then raise ArgumentError, "Must provide a hostname" end
 
 		super(config)
-
+		
 		@user = filter_command(config['user'])
 		@host = filter_command(config['host'])
 	end
@@ -57,16 +56,22 @@ class RemoteEsxDriver < VmDriver
 	end
 
 	def revert_snapshot(snapshot)
-		raise "Not Implemented"
 
-
-
-		#vmware-vim-cmd vmsvc/snapshot.revert [vmid: int] [snapshotlevel: int] [snapshotindex: int]
-		# not sure how we can do this, would have to list snapshots and map name to level & index
-
-		#snapshot = filter_input(snapshot)
-		#system_command("ssh #{@user}@#{@host} \"vim-cmd vmsvc/snapshot.revert #{@vmid} 0 0\"")
+		snapshots = get_snapshots
+		
+		# Look through our snapshot list, choose the right one based on display_name		
+		snapshots.each do |snapshot_obj|
+			if snapshot_obj[:display_name].downcase == snapshot.downcase
+				snapshot_number = snapshot_obj[:name].join(" ")
+				system_command("ssh #{@user}@#{@host} \"vim-cmd vmsvc/snapshot.revert #{@vmid} #{snapshot_number}\"")
+				return true
+			end
+		end
+	
+		# If we got here, the snapshot didn't exist
+		raise "Invalid Snapshot Name"
 	end
+		
 
 	def delete_snapshot(snapshot)
 		raise "Not Implemented"
@@ -113,6 +118,43 @@ class RemoteEsxDriver < VmDriver
 	false
 	end
 
+private 
+
+	def get_snapshots
+		# Command take the format: 
+		# vmware-vim-cmd vmsvc/snapshot.revert [vmid: int] [snapshotlevel: int] [snapshotindex: int]
+		output = `ssh #{@user}@#{@host} \"vim-cmd vmsvc/snapshot.get #{@vmid}\"`
+
+		# this keeps track of the snapshots, takes the form:
+		#[ {:name => [0,0], :display_name => "String containing the snapshotname}, 
+		#  {:name => [0,1], :display_name => "String containing the snapshotname}, ]
+		#  ... 
+		snapshots = []
+		
+		# Use these to keep track of the parsing...		
+		current_tree = -1
+		current_num = 0
+		count = 0
+		
+		# Do the parsing & stick the snapshots in the snapshots array
+		output_lines = output.split("\n")
+		output_lines.each do |line|
+			if line.include?("|") # this is a new snapshot
+				if line.include?("ROOT") # it's a root
+					current_tree = current_tree + 1 # new tree
+					snapshots << { :name => [current_tree,current_num], :display_name => output_lines[count+1].split(":").last.strip }
+					current_num = 0
+				else
+					current_num = current_num + 1 # new snapshot in current tree
+					snapshots << { :name => [current_tree,current_num], :display_name => output_lines[count+1].split(":").last.strip }
+				end
+			end
+			count = count+1
+		end
+
+	snapshots
+	end
+
 end
 
 end 
diff --git a/lib/lab/driver/vm_driver.rb b/lib/lab/driver/vm_driver.rb
index b82dd5e882..1eb505e3b7 100644
--- a/lib/lab/driver/vm_driver.rb
+++ b/lib/lab/driver/vm_driver.rb
@@ -18,12 +18,13 @@ class VmDriver
 	attr_accessor :tools
 	attr_accessor :credentials
 	
-	def initialize(basic_driver_config)
-		@vmid = filter_command(basic_driver_config["vmid"])
-		@location = filter_command(basic_driver_config["location"])
-		@credentials = basic_driver_config["credentials"] || []
-		@tools = filter_input(basic_driver_config["tools"])
-		@os = filter_input(basic_driver_config["os"])
+	def initialize(config)
+	
+		@vmid = filter_command(config["vmid"].to_s)
+		@location = filter_command(config["location"])
+		@credentials = config["credentials"] || []
+		@tools = filter_input(config["tools"])
+		@os = filter_input(config["os"])
 
 		# Currently only implemented for the first set
 		if @credentials.count > 0
@@ -106,7 +107,7 @@ private
 	def scp_to(from,to)
 		require 'net/scp'
 
-		Net::SCP.start(@vmid, @vm_user, :password => @vm_pass) do |scp|
+		Net::SCP.start(@hostname, @vm_user, :password => @vm_pass) do |scp|
 			scp.upload!(from,to)
 		end	
 	end
@@ -115,13 +116,13 @@ private
 		require 'net/scp'
 
 		# download a file from a remote server
-		Net::SCP.start(@vmid, @vm_user, :password => @vm_pass) do |scp|
+		Net::SCP.start(@hostname, @vm_user, :password => @vm_pass) do |scp|
 			scp.download!(from,to)
 		end
 	end
 	
 	def ssh_exec(command)
-		Net::SSH.start(@vmid, @vm_user, :password => @vm_pass) do |ssh|
+		Net::SSH.start(@hostname, @vm_user, :password => @vm_pass) do |ssh|
 			result = ssh.exec!(command)
 		end
 	end
@@ -130,7 +131,7 @@ private
 		return "" unless string # nil becomes empty string
 		return unless string.class == String # Allow other types unmodified
 		
-		if !(string =~ /^[0-9\w\s\[\]\{\}\/\\\.\-\"\(\):!]*$/)
+		unless /^[\w\s\[\]\{\}\/\\\.\-\"\(\):!]*$/.match string
 			raise "WARNING! Invalid character in: #{string}"
 		end
 
@@ -141,7 +142,7 @@ private
 		return "" unless string # nil becomes empty string
 		return unless string.class == String # Allow other types unmodified		
 		
-		if !(string =~ /^[0-9\w\s\[\]\{\}\/\\\.\-\"\(\)]*$/)
+		unless /^[\w\s\[\]\{\}\/\\\.\-\"\(\)]*$/.match string
 			raise "WARNING! Invalid character in: #{string}"
 		end
 
diff --git a/lib/lab/modifier/meterpreter_modifier.rb b/lib/lab/modifier/meterpreter_modifier.rb
index f0deeafc00..8d4fd3b738 100644
--- a/lib/lab/modifier/meterpreter_modifier.rb
+++ b/lib/lab/modifier/meterpreter_modifier.rb
@@ -10,7 +10,7 @@ end
 
 
 # This allows us to override the default way of running commands
-# Currently useful for the esx controller 
+# Currently useful for the esx controller
 
 module Lab
 class Vm
@@ -20,7 +20,6 @@ class Vm
 	attr_accessor :session_input
 	attr_accessor :session_output
 
-
 	def create_framework
 		return if @framework
 		@framework    = Msf::Simple::Framework.create
@@ -42,7 +41,7 @@ class Vm
 		if @os == "windows"
 			exploit_name = 'windows/smb/psexec'
 			payload_name = 'windows/meterpreter/bind_tcp'
-			options = {	"RHOST"		=> @vmid, 
+			options = {	"RHOST"		=> @hostname, 
 					"SMBUser" 	=> @vm_user, 
 					"SMBPass" 	=> @vm_pass}
 
@@ -64,7 +63,7 @@ class Vm
 		else
 			module_name = 'scanner/ssh/ssh_login'
 			payload_name = 'linux/x86/meterpreter/bind_tcp'
-			options = {	"RHOSTS"		=> @vmid, 
+			options = {	"RHOSTS"		=> @hostname, 
 					"USERNAME" 		=> @vm_user, 
 					"PASSWORD" 		=> @vm_pass, 
 					"BLANK_PASSWORDS" 	=> false, 
diff --git a/lib/lab/modifier/test_modifier.rb b/lib/lab/modifier/test_modifier.rb
index d1d7d330c6..0a48e08fcb 100644
--- a/lib/lab/modifier/test_modifier.rb
+++ b/lib/lab/modifier/test_modifier.rb
@@ -4,14 +4,6 @@
 module Lab
 module Modifier
 module Test
-	def dhclient
-		run_command("sudo dhclient")
-	end
-	
-	def route
-		run_command("route -n")
-	end
-
 	def install_nmap
 		run_command("sudo apt-get install nmap")
 	end
diff --git a/lib/lab/modifiers.rb b/lib/lab/modifiers.rb
index 58eb5a7479..e29d341195 100644
--- a/lib/lab/modifiers.rb
+++ b/lib/lab/modifiers.rb
@@ -1,3 +1,4 @@
 require 'modifier/test_modifier'
 require 'modifier/backtrack5_modifier'
+require 'modifier/linux_modifier'
 require 'modifier/meterpreter_modifier'
diff --git a/lib/lab/vm.rb b/lib/lab/vm.rb
index b24497237a..48b5e349cc 100644
--- a/lib/lab/vm.rb
+++ b/lib/lab/vm.rb
@@ -31,11 +31,18 @@ class Vm
 	
 	def initialize(config = {})	
 
-		# TODO - This is such a mess. clean up, and pass stuff down to drivers
+		# TODO - This is a mess. clean up, and pass stuff down to drivers
 		# and then rework the code that uses this api. 
-		@vmid = config['vmid'] 
+		@vmid = config['vmid'].to_s 
 		raise "Invalid VMID" unless @vmid
 
+		# Grab the hostname if specified, otherwise use the vmid
+		# VMID will be different in the case of ESX
+		@hostname = config['hostname']
+		if !@hostname
+			@hostname = @vmid
+		end
+
 		@driver_type = filter_input(config['driver'])
 		@driver_type.downcase!
 
@@ -86,8 +93,7 @@ class Vm
 		else
 			raise "Unknown Driver Type"
 		end
-		
-		
+				
 		# Load in a list of modifiers. These provide additional methods
 		# Currently it is up to the user to verify that 
 		# modifiers are properly used with the correct VM image.
@@ -145,8 +151,8 @@ class Vm
 	end
 
 	def revert_and_start(snapshot)
-		self.revert_snapshot(snapshot)
-		self.start
+		@driver.revert_snapshot(snapshot)
+		@driver.start
 	end
 
 	def copy_to(from,to)
@@ -225,7 +231,7 @@ private
 		return "" unless string # nil becomes empty string
 		return unless string.class == String # Allow other types
 					
-		if !(string =~ /^[(!)\d*\w*\s*\[\]\{\}\/\\\.\-\"\(\)]*$/)
+		unless /^[(!)\d*\w*\s*\[\]\{\}\/\\\.\-\"\(\)]*$/.match string
 			raise "WARNING! Invalid character in: #{string}"
 		end