From f164059d6076e91bd87fbaec5d479f2b2111b906 Mon Sep 17 00:00:00 2001
From: Jonathan Cran <jcran@rapid7.com>
Date: Tue, 28 Dec 2010 19:47:33 +0000
Subject: [PATCH] backend changes for lab functionality. architecture now
 allows storage of more info about VMs, and Labs made up of
 different-technology vms (ie, 1 vmware, 1 vbox, 1 amazon). drivers still need
 to be written

git-svn-id: file:///home/svn/framework3/trunk@11433 4d416f70-5f16-0410-b530-b9f4589650da
---
 lib/lab.rb                    |  11 --
 lib/lab/lab_controller.rb     | 351 ----------------------------------
 lib/lab/vm.rb                 | 124 ++++++++++++
 lib/lab/vm_controller.rb      | 125 ++++++++++++
 lib/lab/vm_driver.rb          |  74 +++++++
 lib/lab/vmware_controller.rb  | 153 ---------------
 lib/lab/workstation_driver.rb |  99 ++++++++++
 7 files changed, 422 insertions(+), 515 deletions(-)
 delete mode 100644 lib/lab.rb
 delete mode 100644 lib/lab/lab_controller.rb
 create mode 100644 lib/lab/vm.rb
 create mode 100644 lib/lab/vm_controller.rb
 create mode 100644 lib/lab/vm_driver.rb
 delete mode 100644 lib/lab/vmware_controller.rb
 create mode 100644 lib/lab/workstation_driver.rb

diff --git a/lib/lab.rb b/lib/lab.rb
deleted file mode 100644
index cd71bf3128..0000000000
--- a/lib/lab.rb
+++ /dev/null
@@ -1,11 +0,0 @@
-##
-## $Id$
-##
-## Requiring this file will give access to lab functionality. Specifically
-##  the ability to start/stop/snapshot/revert & run commands on VMs
-##
-## $Revision$
-##
-
-# Controllers
-require 'lab/lab_controller'
diff --git a/lib/lab/lab_controller.rb b/lib/lab/lab_controller.rb
deleted file mode 100644
index 6302d63e91..0000000000
--- a/lib/lab/lab_controller.rb
+++ /dev/null
@@ -1,351 +0,0 @@
-##
-## $Id$
-##
-## This is the main lab controller which will call out to other controller
-## libraries. Requiring this file and specifying the type of VM at initialization
-## will allow you to start/stop/snapshot/revert & run commands on VMs
-##
-## $Revision$
-##
-
-$:.unshift(File.expand_path(File.dirname(__FILE__))) ## Msf Test libraries
-
-require 'find'
-require 'vmware_controller'
-
-## Not implemented yet!
-#require 'vbox_controller'
-#require 'qemu_controller'
-#require 'ec2_controller'
-
-#
-# ~Higher-level lab methods which are generic to the types of things we want to do with a lab of machines
-#  Note that any generic vm functionality should be pushed down into the controller class. 
-
-class LabController 
-
-	attr_accessor :labdef
-	attr_accessor :controller
-
-	def initialize (labdef = nil, labtype="vmware")
-		
-		if !labdef 
-			## Just use a blank lab to start
-			@labdef = {}
-		else
-			@labdef = labdef
-		end		
-
-		### AHH, traverse again, looking for nils, cuz ruby is teh fail - TODO -  value?		
-		@labdef.each do |key,value|
-			@labdef[key].each do |subkey,subvalue|
-				if !subvalue
-					@labdef[key][subkey] = ""
-				end
-			end
-		end
-		
-				
-		## set up the controller. note that this is likely to change in the future (to provide for additional vm libs)
-		if labtype == "vmware"
-			@controller = VmwareController.new
-			@file_extension = "vmx"
-
-##	
-##	Not implemented yet!
-##
-##		elsif labtype == "qemu"
-##			@controller = QemuController.new
-##			@file_extension = "img"
-##		elsif labtype == "ec2"
-##			@controller = Ec2Controller.new
-
-		else
-			raise "Invalid Lab Controller"
-		end
-	end
-
-	def contains?(vmid)
-		if get_full_path(vmid)
-		 	true
-	 	end
-	end 
-	
-	def build_lab_from_running()
-		vm_array = self.list_lab_running.split("\n") ## this should probably return an array
-		vm_array.shift
-		hlp_stuff_array_into_lab(vm_array)
-	end
-
-	def build_lab_from_files(basepath=nil)
-		vm_array = Find.find(basepath).select { |f| 
-			f =~ /\.#{@file_extension}$/ #&& File.executable?(f)
-		}
-
-		hlp_stuff_array_into_lab(vm_array)
-	end
-
-	def hlp_stuff_array_into_lab(arr)
-		return false unless arr.kind_of? Array
-	
-		arr.each_with_index {|v,i|
-			## give us a vmid!
-			index = @labdef.count + 1
-
-			## eliminate dooops
-			fresh = true
-			@labdef.each { |vmid, definition| fresh = false if labdef[vmid]['file'] == 'v' }
-			
-			if fresh
-				@labdef[index.to_s] = {'file' => v }
-			end
-			#end
-		}
-		return @labdef
-	end
-
-	def run_command_on_lab_vm(vmid,command,arguments=nil)
-		if @labdef[vmid]["tools"]
-			@controller.run_command(get_full_path(vmid), command, @labdef[vmid]["user"],@labdef[vmid]["pass"])
-		else
-			if @labdef[vmid]["os"] == "linux"
-				@controller.run_ssh_command(@labdef[vmid]["hostname"], command , @labdef[vmid]["user"])
-			else
-				raise "OS Not Supported"
-			end
-		end
-	end
-
-	def run_browser_on_lab_vm(vmid,uri)
-		if @labdef[vmid]["os"] == "linux"
-			command = "firefox " + uri
-		elsif @labdef[vmid]['os'] == "windows"
-			command = "C:\\Progra~1\\intern~1\\iexplore.exe " + uri
-		else
-			raise "Don't know how to browse to '" + uri + "'."
-		end
-		run_command_on_lab_vm(vmid,command)
-	end
-
-	def start_lab_vm(vmid)
-		if running?(vmid)
-			self.list_lab_running
-		else
-			@controller.start(get_full_path(vmid)) 
-		end
-	end
-	
-	def reset_lab_vm(vmid)
-		@controller.reset(get_full_path(vmid))
-	end
-
-	def pause_lab_vm(vmid)
-		if !running?(vmid)
-			self.list_lab_running
-		else
-			@controller.pause(get_full_path(vmid)) 
-		end
-	end
-
-	def suspend_lab_vm(vmid)
-		if !running?(vmid)
-			self.list_lab_running
-		else
-				@controller.suspend(get_full_path(vmid))
-		end
-	end
-
-	def snapshot_lab_vm(vmid, snapshot)
-		if !running?(vmid)
-			self.list_lab_running
-		else
-			@controller.create_snapshot(get_full_path(vmid),snapshot)
-		end
-	end
-
-	def revert_lab_vm(vmid, snapshot, run=true)
-		if !running?(vmid)
-			self.list_lab_running
-		else
-			@controller.revert_snapshot(get_full_path(vmid),snapshot)
-
-			## If you revert w/ vmrun, you need to restart
-			if run 
-				sleep 5
-				@controller.start(get_full_path(vmid))
-			end
-		end
-	end
-
-	def stop_lab_vm(vmid)
-		if !running?(vmid)
-			self.list_lab_running
-		else
-			@controller.stop(get_full_path(vmid))
-		end
-	end
-
-	def start_lab
-		@labdef.each { | key, value |
-			if value
-				self.start_lab_vm(key)
-			end
-		}
-	end
-
-	def pause_lab
-		@labdef.each { | key, value |
-			if value
-				self.pause_lab_vm(key)
-			end
-
-		}
-	end
-
-
-	def stop_lab
-		@labdef.each { | key, value |
-			if value
-				self.stop_lab_vm(key)
-			end
-		}
-
-	end
-
-	def reset_lab
-		@labdef.each { | key, value |
-			if value
-				self.reset_lab_vm(key)
-			end
-		}
-	end
-
-
-	def suspend_lab
-		@labdef.each { | key, value |
-			if value 
-				self.suspend_lab_vm(key)
-			end
-		}
-	end
-
-	def snapshot_lab(snapshot)
-		@labdef.each { | key, value |
-			if value
-				self.snapshot_lab_vm(key,snapshot)
-			end
-		}
-	end
-
-	def revert_lab(snapshot)
-		@labdef.each { | key, value |
-			if value
-				self.revert_lab_vm(key,snapshot)
-			end
-		}
-	end
-
-        def run_command_on_lab(command)
-                @labdef.each { | key, value |
-                        if value
-                                self.run_command_on_lab_vm(key, command)
-                        end
-                }
-        end
-
-	def copy_to_lab(file)	
-		@labdef.each { | key, value |
-			if value
-				self.copy_to_lab_vm(key,file)
-			end
-		}
-	end
-
-	def copy_from_lab(file)	
-		@labdef.each { | key, value | 
-#			next unless line =~ /^#{@vmbase}(.*)/
-			if value
-				self.copy_from_lab_vm(key,file)
-			end
-		}
-	end
-
-	def copy_to_lab_vm(vmid,file)
-		## handle linux
-		
-		guestpath  = ""
-		
-		if (@labdef[vmid]["os"] == "linux")
-			guestpath = "/tmp/"
-		else
-			guestpath = "C:\\temp_msf\\\\" ## double-escaping because it's being used in a system command. 
-		end
-
-		name = File.basename(file)
-
-		## if we've installed vm-tools on the box, use that to copy. if not, use scp
-		if (@labdef[vmid]["tools"] == "true")			
-			@controller.create_directory_in_guest(get_full_path(vmid),@labdef[vmid]["user"],@labdef[vmid]["pass"], guestpath)
-			@controller.copy_file_to(get_full_path(vmid),@labdef[vmid]["user"],@labdef[vmid]["pass"],file, guestpath + name)
-		else
-			@controller.scp_copy_file_to(@labdef[vmid]["hostname"], @labdef[vmid]["user"], file, guestpath + name)
-		end
-	end
-	
-	def copy_from_lab_vm(vmid,file)
-		hostpath = "/tmp/"
-
-		name = File.basename(file.gsub("\\","/")) 
-
-		@controller.copy_file_from(get_full_path(vmid),@labdef[vmid]["user"],@labdef[vmid]["pass"],file,hostpath + name)
-	end
-	
-	def list_lab_running
-		@controller.get_running	
-	end
-
-	def list_lab
-		str = ""
-		@labdef.sort.each { |key,val|
-
-			if val != nil
-				str = str + key.to_s + ": "  + val["file"].to_s + "\n"
-			end
-		 }
-		return str
-	end
-
-        def find_lab_vm(search)
-		str = ""
-                @labdef.sort.each { |key,val|
-                        if val != nil
-				if (val["file"].to_s.downcase.index(search.downcase) != nil)
-	                                str = str + key.to_s + ": "  + val["file"].to_s + "\n"
-				end
-                        end
-                 }
-		return str
-        end
-
-	def exists?(vmid)
-		if get_full_path(vmid)
-			return true
-		end
-		return false
-	end
-
-	def running?(vmid)
-		if @controller.running?(get_full_path(vmid))
-			return true
-		end
-		return false
-	end
-
-	def get_full_path(vmid)
-		if @labdef[vmid]
-			@labdef[vmid]["file"].to_s		## handle linux
-		else
-			nil
-		end
-	end
-end
diff --git a/lib/lab/vm.rb b/lib/lab/vm.rb
new file mode 100644
index 0000000000..ff2add3768
--- /dev/null
+++ b/lib/lab/vm.rb
@@ -0,0 +1,124 @@
+##
+## $Id$
+##
+
+require 'workstation_driver'
+#require 'server_driver'
+#require 'qemu_driver'
+#require 'vbox_driver'
+#require 'ec2_driver'
+#require 'azure_driver'
+
+class Vm
+	
+	attr_accessor :vmid
+	attr_accessor :driver
+	attr_accessor :location
+	attr_accessor :credentials
+	attr_accessor :tools
+
+	## Initialize takes a vm configuration hash of the form
+	##  - driver (vm technology)
+	##  - vmid (unique identifier)
+	##  - location (file / uri)
+	##  - credentials (of the form [ {'user'=>"user",'pass'=>"pass", 'admin' => false}, ... ])
+	def initialize(config = {})	
+		@driver = nil
+		driver_type = config['driver']
+		driver_type.downcase!
+
+		@vmid = config['vmid']
+		@location = config['location']
+
+		## Internals
+		@credentials = config['credentials'] || []
+		@tools = config['tools'] || false		## TODO
+		@operating_system = nil				## TODO
+		@ports = nil					## TODO
+		@vulns = nil					## TODO
+
+		if driver_type == "workstation"
+			@driver = WorkstationDriver.new(@location)
+		#elsif driver_type == "server"
+		#	@driver = ServerDriver.new
+		#elsif driver_type == "virtualbox"
+		#	@driver = VirtualBoxDriver.new	
+		#elsif driver_type == "qemu"
+		#	@driver = QemuDriver.new	
+		#elsif driver_type == "ec2"
+		#	@driver = Ec2Driver.new	
+		#elsif driver_type == "azure"
+		#	@driver = AzureDriver.new	
+		else
+			raise Exception, "Unknown Driver Type"
+		end
+	end
+
+	
+	def running?
+		@driver.running?
+	end
+
+	def start
+		@driver.start
+	end
+
+	def stop
+		@driver.stop
+	end
+
+	def pause
+		@driver.pause
+	end
+
+	def suspend
+		@driver.suspend
+	end
+	
+	def resume
+		@driver.resume
+	end
+
+	def snapshot(name)
+		@driver.snapshot(name)
+	end
+
+	def revert(name)
+		@driver.revert(name)
+	end
+
+	def copy_to(from_file,to_file)
+		raise Exception, "not implemented"
+	end
+	
+	def copy_from(from_file,to_file)
+		raise Exception, "not implemented"
+	end
+	
+	def run_command(command,arguments=nil)
+		raise Exception, "not implemented"
+	end
+
+	def open_uri(uri)
+		raise Exception, "not implemented"
+	end
+
+	def to_s
+		return @vmid.to_s + ": " + @location.to_s
+	end
+
+	def to_yaml
+		out =  " - vmid: #{vmid}\n"
+		out += "   driver: #{driver.type}\n"
+		out += "   location: #{location}\n"
+		out += "   tools: #{tools}\n"
+		out += "   credentials:\n"
+		@credentials.each do |credential|		
+			out += "     - user: #{credential['user']}\n"
+			out += "       pass: #{credential['pass']}\n"
+			out += "       admin: #{credential['admin']}\n"
+		end
+
+	 	return out
+	end		
+end
diff --git a/lib/lab/vm_controller.rb b/lib/lab/vm_controller.rb
new file mode 100644
index 0000000000..f77a2ccc3b
--- /dev/null
+++ b/lib/lab/vm_controller.rb
@@ -0,0 +1,125 @@
+##
+## $Id$
+##
+## This is the main lab controller which will call out to other controller
+## libraries. Requiring this file and specifying the type of VM at initialization
+## will allow you to start/stop/snapshot/revert & run commands on VMs
+##
+## $Revision$
+##
+
+$:.unshift(File.expand_path(File.dirname(__FILE__))) ## Msf Test libraries
+
+require 'find'
+require 'enumerator'
+require 'vm'
+require 'yaml'
+
+#
+# ~Higher-level lab methods which are generic to the types of things we want to do with a lab of machines
+#  Note that any generic vm functionality should be pushed down into the controller class. 
+
+class VmController 
+
+	include Enumerable
+
+	def initialize (labdef = nil)
+		
+		@vms = [] ## Start with an empty array of vms
+
+		## labdef is a big array of hashes (vms) - generally loaded from yaml
+		if !labdef 
+			labdef = [] ## Just use a blank lab to start
+		else
+			labdef = labdef
+		end
+
+		## Create vm objects from the lab definition
+		labdef.each do |item|
+			begin
+				@vms << Vm.new(item)
+			rescue Exception => e
+				puts e.to_s
+			end  
+		end
+
+	end
+	
+	def clear!
+		@vms = []
+	end
+
+	def find_by_vmid(vmid)
+		@vms.each do |vm|
+
+			if (vm.vmid.to_s == vmid.to_s)
+				return vm
+			end
+		end
+		return nil
+	end
+
+	def from_file(file)
+		labdef = YAML::load_file(file)
+
+		labdef.each do |item|
+			puts "Lab item: " + item.inspect
+			@vms << Vm.new(item)
+		end
+	end
+
+	def to_file(file)
+		File.open(file, 'w') do |f|
+			@vms.each { |vm| f.puts vm.to_yaml }
+		end
+	end
+
+	def each
+		@vms.each { |vm| yield vm }
+	end
+
+#	def includes?(vm)
+#		@vms.each { |vm| if (vm.vmid.to_ == vmid.to_s) then return true end  }
+#	end
+
+	def includes_vmid?(vmid)
+		@vms.each { |vm| if (vm.vmid.to_s == vmid.to_s) then return true end  }
+	end
+
+	def running?(vmid)
+		if exists?(vmid)
+			return self.find_by_vmid(vmid).running?
+		end
+		return false 
+	end
+
+	## Might want to mix this (workstation) functionality in?
+	def build_from_running_workstation(clear=false)
+
+		if clear
+			@vms = []
+		end
+		
+		vm_list = `vmrun list`.split("\n")
+		vm_list.shift
+		vm_list.each do |vmx|
+			index = @vms.count + 1 ## give us a vmid!
+			@vms << Vm.new( {"vmid" => index, "driver" => "workstation", 
+					"location" => vmx})
+		end
+	end
+
+	def build_from_dir_workstation(basepath=nil, clear=false)
+	
+		if clear
+			@vms = []
+		end
+
+		vm_list = Find.find(basepath).select { |f| f =~ /\.vmx$/ }
+		vm_list.each do |vmx|
+			index = @vms.count + 1 ## give us a vmid!
+			@vms << Vm.new( {"vmid" => index, "driver" => "workstation",
+					"location" => vmx})
+		end
+	end
+end
diff --git a/lib/lab/vm_driver.rb b/lib/lab/vm_driver.rb
new file mode 100644
index 0000000000..c6183d508e
--- /dev/null
+++ b/lib/lab/vm_driver.rb
@@ -0,0 +1,74 @@
+##
+## $Id$
+##
+
+class VmDriver
+
+	attr_accessor :type
+	attr_accessor :location
+
+	def initialize(location)
+	end
+
+	def start
+	end
+
+	def stop
+	end
+
+	def suspend
+	end
+
+	def pause
+	end
+
+	def reset
+	end
+
+	def snapshot(snapshot)
+	end
+
+	def revert(snapshot)
+	end
+
+	def delete_snapshot(snapshot)
+	end
+
+	def run_command(command, user, pass)	
+	end
+	
+	def copy_from(user, pass, from, to)
+	end
+	
+	def copy_to(user, pass, from, to)
+	end
+
+	def check_file_exists(user, pass, file)
+	end
+
+	def create_directory(user, pass, directory)
+	end
+
+	def ssh_exec(host, command, user)
+		ssh_command = "ssh " + user + "@" + host + " " + command
+		system_command(ssh_command)
+	end
+
+	def scp_from(host, user, from, to)
+		vmrunstr = "scp -r \"" + user + "@" + host + ":" + from + "\" \"" + to + "\"" ## TODO - setup keys  
+		system_command(vmrunstr)
+	end
+
+	def scp_to(host, user, from, to)
+		vmrunstr = "scp -r \"" + from + "\" \"" + user + "@" + host + ":" + to + "\"" ## TODO - setup keys  
+		system_command(vmrunstr)
+	end
+
+	private
+
+		def system_command(command)
+			## TODO - filter here
+			system(command)
+		end
+
+end
diff --git a/lib/lab/vmware_controller.rb b/lib/lab/vmware_controller.rb
deleted file mode 100644
index 934e369284..0000000000
--- a/lib/lab/vmware_controller.rb
+++ /dev/null
@@ -1,153 +0,0 @@
-##
-## $Id$
-##
-## This class contains lower-level methods (than the lab_controller) and is specific
-## to vmware. Do not require this file directly. Instead, require the lab_controller, 
-## and specify the type of VM at initialization. 
-##
-## $Revision$
-##
-
-class VmwareController
-
-	def initialize
-	end
-
-	def start(vmx)
-		if File.exist?(vmx)
-			system_command("vmrun -T ws start " + "\"" + vmx + "\"")
-		else
-			raise ArgumentError, "Couldn't find: " + vmx, caller
-		end
-	end
-
-	def stop(vmx)
-		if File.exist?(vmx)
-			system_command("vmrun -T ws stop " + "\"" + vmx + "\"")
-		else
-			raise ArgumentError,"Couldn't find:  " + vmx, caller
-		end
-	end
-
-	def suspend(vmx)
-		if File.exist?(vmx)
-			system_command("vmrun -T ws suspend " + "\"" + vmx + "\"")
-		else
-			raise ArgumentError,"Couldn't find: " + vmx, caller
-		end
-	end
-
-	def pause(vmx)
-		if File.exist?(vmx)
-			system_command("vmrun -T ws pause " + "\"" + vmx + "\"")
-		else
-			raise ArgumentError, "Couldn't find: " + vmx, caller
-		end
-	end
-
-	def reset(vmx)
-		if File.exist?(vmx)
-			system_command("vmrun -T ws reset " + "\"" + vmx + "\"")
-		else
-			raise ArgumentError, "Couldn't find: " + vmx, caller
-		end
-	end
-
-	def run_command(vmx, command, user, pass, displayParameter=false)
-		if File.exist?(vmx)
-			
-			vmrunstr = "vmrun -T ws -gu \"" + user + "\" -gp \"" + pass + "\" runProgramInGuest \"" + vmx + "\" " + "\"" + command + "\" -noWait -activeWindow"
-	
-			if displayParameter
-				vmrunstr = vmrunstr + " -display :0"
-			end
-
-			system_command(vmrunstr)
-		else
-			raise ArgumentError,"Couldn't find: " + vmx, caller
-		end
-	end
-	
-	def run_ssh_command(hostname, command, user)
-		ssh_command = "ssh " + user + "@" + hostname + " " + command
-		system_command(ssh_command)
-	end
-
-	def copy_file_from(vmx, user, pass, guestpath, hostpath)
-		vmrunstr = "vmrun -T ws -gu " + user + " -gp " + pass + " copyFileFromGuestToHost \"" + vmx + "\" \"" + guestpath + "\" \"" + hostpath + "\"" 
-		system_command(vmrunstr)
-	end
-	
-	def scp_copy_file_from(hostname, user, guestpath, hostpath)
-		vmrunstr = "scp -r \"" + user + "@" + hostname + ":" + guestpath + "\" \"" + hostpath + "\"" ## TODO - setup keys  
-		system_command(vmrunstr)
-	end
-
-	def copy_file_to(vmx, user, pass, hostpath, guestpath)
-		vmrunstr = "vmrun -T ws -gu " + user + " -gp " + pass + " copyFileFromHostToGuest \"" + vmx + "\" \"" + hostpath + "\" \"" + guestpath + "\""  
-		system_command(vmrunstr)
-	end
-
-	def scp_copy_file_to(hostname, user, hostpath, guestpath)
-		vmrunstr = "scp -r \"" + hostpath + "\" \"" + user + "@" + hostname + ":" + guestpath + "\"" ## TODO - setup keys  
-		system_command(vmrunstr)
-	end
-
-	def check_file_exists(vmx, user, pass, file)
-		vmrunstr = "vmrun -T ws -gu " + user + " -gp " + pass + " fileExistsInGuest \"" + vmx + "\" \"" + file + "\" "
-		system_command(vmrunstr)
-	end
-
-	def create_directory_in_guest(vmx, user, pass, directory)
-		vmrunstr = "vmrun -T ws -gu " + user + " -gp " + pass + " createDirectoryInGuest \"" + vmx + "\" \"" + directory + "\" "
-		system_command(vmrunstr)
-	end
-
-	def create_snapshot(vmx, snapshot)
-		if File.exist?(vmx)
-			system_command("vmrun -T ws snapshot " + "\"" + vmx + "\" " + snapshot)
-		else
-			raise ArgumentError,"Couldn't find: " + vmx, caller
-		end
-	end
-
-	def revert_snapshot(vmx, snapshot)
-		if File.exist?(vmx)
-			system_command("vmrun -T ws revertToSnapshot " + "\"" + vmx + "\" " + snapshot)
-		else
-			raise "Couldn't find: " + vmx, caller
-		end
-	end
-
-	def delete_snapshot(vmx, snapshot)
-		if File.exist?(vmx)
-			system_command("vmrun -T ws deleteSnapshot " + "\"" + vmx + "\" " + snapshot )
-		else
-			raise ArgumentError,"Couldn't find: " + vmx, caller
-		end
-	end
-
-	def get_running
-		output = `vmrun list` ##hackity hack=begin
-		return output
-	end
-
-	def running?(vmx)
-		output = self.get_running
-		output.each_line do |line| 
-			if line.strip == vmx.strip
-				return true
-			end
-		end
-
-		return false
-	end
-
-	private
-
-	def system_command(command)
-		puts "DEBUG: " + command
-		system(command)
-	end
-
-end
diff --git a/lib/lab/workstation_driver.rb b/lib/lab/workstation_driver.rb
new file mode 100644
index 0000000000..38870e4361
--- /dev/null
+++ b/lib/lab/workstation_driver.rb
@@ -0,0 +1,99 @@
+require 'vm_driver'
+
+##
+## $Id$
+##
+
+class WorkstationDriver < VmDriver
+
+	attr_accessor :type
+	attr_accessor :location
+
+	def initialize(location)
+		if !File.exist?(location)
+			raise ArgumentError,"Couldn't find: " + location
+		end
+
+		@location = location
+		@type = "Workstation"
+	end
+
+	def start
+		system_command("vmrun -T ws start " + "\"" + @location + "\"")
+	end
+
+	def stop
+		system_command("vmrun -T ws stop " + "\"" + @location + "\"")
+	end
+
+	def suspend
+		system_command("vmrun -T ws suspend " + "\"" + @location + "\"")
+	end
+
+	def pause
+		system_command("vmrun -T ws pause " + "\"" + @location + "\"")
+	end
+
+	def reset
+		system_command("vmrun -T ws reset " + "\"" + @location + "\"")
+	end
+
+	def create(snapshot)
+		system_command("vmrun -T ws snapshot " + "\"" + @location + "\" " + snapshot)
+	end
+
+	def revert(snapshot)
+		system_command("vmrun -T ws revertToSnapshot " + "\"" + @location + "\" " + snapshot)
+	end
+
+	def delete_snapshot(snapshot)
+		system_command("vmrun -T ws deleteSnapshot " + "\"" + @location + "\" " + snapshot )
+	end
+
+	def run_command(command, user, pass)
+		vmrunstr = "vmrun -T ws -gu \"" + user + "\" -gp \"" + pass + "\" runProgramInGuest \"" + 
+				@location + "\" " + "\"" + command + "\" -noWait -activeWindow"
+
+		system_command(vmrunstr)
+	end
+	
+	def copy_from(user, pass, from, to)
+		vmrunstr = "vmrun -T ws -gu " + user + " -gp " + pass + " copyFileFromGuestToHost \"" +
+				@location + "\" \"" + from + "\" \"" + to + "\"" 
+		system_command(vmrunstr)
+	end
+
+	def copy_to(user, pass, from, to)
+		vmrunstr = "vmrun -T ws -gu " + user + " -gp " + pass + " copyFileFromHostToGuest \"" + 
+				@location + "\" \"" + from + "\" \"" + to + "\""  
+		system_command(vmrunstr)
+	end
+
+	def check_file_exists(user, pass, file)
+		vmrunstr = "vmrun -T ws -gu " + user + " -gp " + pass + " fileExistsInGuest \"" +
+				@location + "\" \"" + file + "\" "
+		system_command(vmrunstr)
+	end
+
+	def create_directory(user, pass, directory)
+		vmrunstr = "vmrun -T ws -gu " + user + " -gp " + pass + " createDirectoryInGuest \"" + 
+				@location + "\" \"" + directory + "\" "
+		system_command(vmrunstr)
+	end
+
+	def running?
+		## Get running Vms
+		running = `vmrun list`
+		running_array = running.split("\n")
+		running_array.shift
+
+		running_array.each do |vmx|
+			if vmx.to_s == @location.to_s
+				return true
+			end
+		end
+
+		return false
+	end
+
+end