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

@@ -13,8 +13,6 @@ class FogDriver < VmDriver
 		super(config)
 		@fog_config = fog_config
 
-		puts "Fog Config: #{fog_config}"
-
 		# Soft dependency
 		begin
 			require 'fog'

lib/lab/driver/remote_esx_driver.rb

@@ -13,7 +13,6 @@ 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
 
@@ -57,16 +56,22 @@ class RemoteEsxDriver < VmDriver
 	end
 
 	def revert_snapshot(snapshot)
-		raise "Not Implemented"
 
+		snapshots = get_snapshots
 		
-
+		# Look through our snapshot list, choose the right one based on display_name		
-		#vmware-vim-cmd vmsvc/snapshot.revert [vmid: int] [snapshotlevel: int] [snapshotindex: int]
+		snapshots.each do |snapshot_obj|
-		# not sure how we can do this, would have to list snapshots and map name to level & index
+			if snapshot_obj[:display_name].downcase == snapshot.downcase
-
+				snapshot_number = snapshot_obj[:name].join(" ")
-		#snapshot = filter_input(snapshot)
+				system_command("ssh #{@user}@#{@host} \"vim-cmd vmsvc/snapshot.revert #{@vmid} #{snapshot_number}\"")
-		#system_command("ssh #{@user}@#{@host} \"vim-cmd vmsvc/snapshot.revert #{@vmid} 0 0\"")
+				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 

lib/lab/driver/vm_driver.rb

@@ -18,12 +18,13 @@ class VmDriver
 	attr_accessor :tools
 	attr_accessor :credentials
 	
-	def initialize(basic_driver_config)
+	def initialize(config)
-		@vmid = filter_command(basic_driver_config["vmid"])
+	
-		@location = filter_command(basic_driver_config["location"])
+		@vmid = filter_command(config["vmid"].to_s)
-		@credentials = basic_driver_config["credentials"] || []
+		@location = filter_command(config["location"])
-		@tools = filter_input(basic_driver_config["tools"])
+		@credentials = config["credentials"] || []
-		@os = filter_input(basic_driver_config["os"])
+		@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
 

lib/lab/modifier/meterpreter_modifier.rb

@@ -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, 

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

lib/lab/modifiers.rb

@@ -1,3 +1,4 @@
 require 'modifier/test_modifier'
 require 'modifier/backtrack5_modifier'
+require 'modifier/linux_modifier'
 require 'modifier/meterpreter_modifier'

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!
 
@@ -87,7 +94,6 @@ class Vm
 			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)
+		@driver.revert_snapshot(snapshot)
-		self.start
+		@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