From 60c7268061ea22584d99b37fe7736cf06e719f47 Mon Sep 17 00:00:00 2001
From: Jonathan Cran <jcran@0x0e.org>
Date: Tue, 13 Mar 2012 23:26:51 -0500
Subject: [PATCH] updated workstation driver

---
 lib/lab/driver/vm_driver.rb          |  6 ++--
 lib/lab/driver/workstation_driver.rb | 46 +++++++++++++---------------
 2 files changed, 24 insertions(+), 28 deletions(-)

diff --git a/lib/lab/driver/vm_driver.rb b/lib/lab/driver/vm_driver.rb
index 82d8d34..178a982 100644
--- a/lib/lab/driver/vm_driver.rb
+++ b/lib/lab/driver/vm_driver.rb
@@ -127,7 +127,7 @@ def remote_system_command(command)
 
 	def scp_to(local,remote)
 		if @vm_keyfile
-			puts "DEBUG: authenticating to #{@hostname} as #{@vm_user} with key #{@vm_keyfile}"
+			#puts "DEBUG: authenticating to #{@hostname} as #{@vm_user} with key #{@vm_keyfile}"
 			Net::SCP.start(@hostname, @vm_user, :keys => [@vm_keyfile]) do |scp|
 				puts "DEBUG: uploading #{local} to #{remote}"
 				scp.upload!(local,remote)
@@ -143,7 +143,7 @@ def scp_to(local,remote)
 	def scp_from(remote, local)
 		# download a file from a remote server
 		if @vm_keyfile
-			puts "DEBUG: authenticating to #{@hostname} as #{@vm_user} with key #{@vm_keyfile}"
+			#puts "DEBUG: authenticating to #{@hostname} as #{@vm_user} with key #{@vm_keyfile}"
 			Net::SCP.start(@hostname, @vm_user, :keys => [@vm_keyfile]) do |scp|
 				puts "DEBUG: downloading #{remote} to #{local}"
 				scp.download!(remote,local)
@@ -158,7 +158,7 @@ def scp_from(remote, local)
 	
 	def ssh_exec(command)
 		if @vm_keyfile
-			puts "DEBUG: authenticating to #{@hostname} as #{@vm_user} with key #{@vm_keyfile}"
+			#puts "DEBUG: authenticating to #{@hostname} as #{@vm_user} with key #{@vm_keyfile}"
 			Net::SSH.start(@hostname, @vm_user, :keys => [@vm_keyfile]) do |ssh|
 				puts "DEBUG: running command: #{command}"
 				ssh.exec!(command)
diff --git a/lib/lab/driver/workstation_driver.rb b/lib/lab/driver/workstation_driver.rb
index fa86df9..0910874 100644
--- a/lib/lab/driver/workstation_driver.rb
+++ b/lib/lab/driver/workstation_driver.rb
@@ -60,27 +60,30 @@ def run_command(command)
 		script_rand_name = rand(1000000)
 
 		#
-		# Configure paths for each OS
+		# Configure paths for each OS - We really can't filter command, so we're gonna 
+		# stick it in a script
 		#
 		if @os == "windows"
 			local_tempfile_path = "/tmp/lab_script_#{script_rand_name}.bat"
 			remote_tempfile_path = "C:\\\\lab_script_#{script_rand_name}.bat"
+			remote_output_file = "C:\\\\lab_command_output_#{script_rand_name}"
 			remote_run_command = remote_tempfile_path
 			File.open(local_tempfile_path, 'w') {|f| f.write(command) }
 		else
+			
 			local_tempfile_path = "/tmp/lab_script_#{script_rand_name}.sh"
 			remote_tempfile_path = "/tmp/lab_script_#{script_rand_name}.sh"
-			remote_run_command = remote_tempfile_path # TODO - do we need to append /bin/sh ?
+			remote_output_file = "/tmp/lab_command_output_#{script_rand_name}"
+			local_output_file = "/tmp/lab_command_output_#{script_rand_name}"
+			
+			remote_run_command = remote_tempfile_path
+			
 			File.open(local_tempfile_path, 'w') {|f| f.write("#!/bin/sh\n#{command}\n")}
 		end
-
-		#
-		# We really can't filter command, so we're gonna stick it in a script
-		#
 		
 		if @tools
 
-			puts "DEBUG: Running w/ tools"
+			#puts "DEBUG: Running w/ tools"
 
 			#
 			# Copy our local tempfile to the guest
@@ -92,19 +95,20 @@ def run_command(command)
 
 			if @os == "linux"
 				#
-				# Now run the command directly on the guest
+				# Now run the command directly on the guest (linux - call w/ /bin/sh)
 				#
 				vmrunstr = "vmrun -T ws -gu #{@vm_user} -gp #{@vm_pass} " + 
-						"runProgramInGuest \'#{@location}\' /bin/sh #{remote_tempfile_path}"
+						"runProgramInGuest \'#{@location}\' /bin/sh #{remote_tempfile_path} > #{remote_output_file}"
 				system_command(vmrunstr)
 			else
 				#
-				# Now run the command directly on the guest
+				# Now run the command directly on the guest (windows)
 				#
 				vmrunstr = "vmrun -T ws -gu #{@vm_user} -gp #{@vm_pass} " + 
-						"runProgramInGuest \'#{@location}\' #{remote_tempfile_path}"
+						"runProgramInGuest \'#{@location}\' #{remote_tempfile_path} > #{remote_output_file}" 
 				system_command(vmrunstr)
 			end
+			
 			#
 			# Cleanup. Delete it on the guest
 			#
@@ -124,11 +128,6 @@ def run_command(command)
 			#
 
 			if @os == "linux"
-				#
-				# Set up our paths
-				#
-				remote_output_file = "/tmp/lab_command_output_#{rand(1000000)}"
-				local_output_file = "/tmp/lab_command_output_#{rand(1000000)}"
 				
 				#
 				# Copy it over
@@ -145,18 +144,15 @@ def run_command(command)
 				#
 				scp_from(remote_output_file, local_output_file)
 
-				#
-				# And clean up after yourself on the remote system
-				#
-				ssh_exec("rm #{remote_output_file}")
-				ssh_exec("rm #{remote_tempfile_path}")
-
 				# Now, let's look at the output of the command
-				string = File.open(local_output_file,"r").read
+				output_string = File.open(local_output_file,"r").read
 
 				#
-				# And clean that up too
+				# And clean up
 				#
+				ssh_exec("rm #{remote_output_file}")
+				ssh_exec("rm #{remote_tempfile_path}")
+				
 				`rm #{local_output_file}`
 
 			else
@@ -164,7 +160,7 @@ def run_command(command)
 			end
 			
 		end
-	return string
+	output_string
 	end
 	
 	def copy_from(from, to)