ssh on workstation

lab.gemspec

@@ -21,7 +21,15 @@ Gem::Specification.new do |s|
 
   # specify any dependencies here; for example:
   # s.add_development_dependency "rspec"
-   s.add_runtime_dependency "nokogiri"
-   s.add_runtime_dependency "net-ssh"
-   s.add_runtime_dependency "net-scp"
+
+    # ??
+    s.add_runtime_dependency "nokogiri"
+
+    # Multiple things - fallback execute / copy
+    s.add_runtime_dependency "net-ssh"
+    s.add_runtime_dependency "net-scp"
+
+    # Vmware ESX driver
+    s.add_runtime_dependency "rbvmomi"
+
 end

lib/lab/driver/remote_workstation_driver.rb

@@ -125,11 +125,10 @@ def copy_from(from, to)
 		system_command(remote_copy_command)
 
 		if @tools 
-
 			remote_system_command("ssh #{@user}@#{@host} \"vmrun -T ws -gu #{@vm_user} -gp #{@vm_pass} " +
 					"copyFileFromGuestToHost \'#{@location}\' \'#{from}\' \'#{to}\' nogui")
 		else
-			scp_to(from,to)
+			scp_from(to,from)
 		end
 	end
 

lib/lab/driver/vm_driver.rb

@@ -5,25 +5,26 @@
 #
 # 	!!WARNING!! - All drivers are expected to filter input before running
 #	anything based on it. This is particularly important in the case
-#	of the drivers which wrap a command line to provide functionality.  
+#	of the drivers which wrap a command line to provide functionality.	
 #
 
 module Lab
 module Drivers
 class VmDriver
-	
+
 	attr_accessor :vmid 
 	attr_accessor :location
 	attr_accessor :os
 	attr_accessor :tools
 	attr_accessor :credentials
 
 	def initialize(config)
-	
+
 		@vmid = filter_command(config["vmid"].to_s)
 		@location = filter_command(config["location"])
 		@credentials = config["credentials"] || []
 		@tools = filter_input(config["tools"])
+		@arch = filter_input(config["arch"])
 		@os = filter_input(config["os"])
 		@hostname = filter_input(config["hostname"]) || filter_input(config["vmid"].to_s)
 
@@ -125,24 +126,47 @@ def remote_system_command(command)
 	end
 
 	def scp_to(local,remote)
-		puts "DEBUG: copying #{local} to #{remote}"
-		Net::SCP.start(@hostname, @vm_user, :password => @vm_pass) do |scp|
-			scp.upload!(from,to)
+		if @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)
+			end
+		else
+			Net::SCP.start(@hostname, @vm_user, :password => @vm_pass, :auth_methods => ["password"]) do |scp|
+				puts "DEBUG: uploading #{local} to #{remote}"
+				scp.upload!(local,remote)
+			end
 		end	
 	end
 
-	def scp_from(local,remote)
+	def scp_from(remote, local)
 		# download a file from a remote server
-		puts "DEBUG: copying #{remote} to #{local}"
-		Net::SCP.start(@hostname, @vm_user, :password => @vm_pass) do |scp|
-			scp.download!(from,to)
+		if @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)
+			end
+		else
+			Net::SCP.start(@hostname, @vm_user, :password => @vm_pass, :auth_methods => ["password"]) do |scp|
+				puts "DEBUG: downloading #{remote} to #{local}"
+				scp.download!(remote,local)
+			end
 		end
 	end
 
 	def ssh_exec(command)
-		puts "DEBUG: ssh command: #{command} on host #{hostname}"
-		Net::SSH.start(@hostname, @vm_user, :password => @vm_pass) do |ssh|
-			result = ssh.exec!(command)
+		if @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)
+			end
+		else
+			Net::SSH.start(@hostname, @vm_user, :password => @vm_pass, :auth_methods => ["password"]) do |ssh|
+				result = ssh.exec!(command)
+			end
 		end
 	end
 
@@ -153,7 +177,6 @@ def filter_input(string)
 		unless /^[\d\w\s\[\]\{\}\/\\\.\-\"\(\):!]*$/.match string
 			raise "WARNING! Invalid character in: #{string}"
 		end
-
 	string
 	end
 
@@ -164,10 +187,9 @@ def filter_command(string)
 		unless /^[\d\w\s\[\]\{\}\/\\\.\-\"\(\)]*$/.match string
 			raise "WARNING! Invalid character in: #{string}"
 		end
-
 	string
 	end
-end
 
 end
 end
+end

lib/lab/driver/workstation_driver.rb

@@ -54,63 +54,115 @@ def delete_snapshot(snapshot)
 
 	def run_command(command)
 
-		script_rand_name = rand(10000)
-
+		#
+		# Generate a script name
+		#
+		script_rand_name = rand(1000000)
+
+		#
+		# Configure paths for each OS
+		#
 		if @os == "windows"
 			local_tempfile_path = "/tmp/lab_script_#{script_rand_name}.bat"
 			remote_tempfile_path = "C:\\\\lab_script_#{script_rand_name}.bat"
 			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 = "/bin/sh #{remote_tempfile_path}"
+			remote_run_command = remote_tempfile_path # TODO - do we need to append /bin/sh ?
+			File.open(local_tempfile_path, 'w') {|f| f.write("#!/bin/sh\n#{command}\n")}
 		end
 
-		# write out our script locally
-		File.open(local_tempfile_path, 'w') {|f| f.write(command) }
-
-		# we really can't filter command, so we're gonna stick it in a script
+		#
+		# We really can't filter command, so we're gonna stick it in a script
+		#
+
 		if @tools
-			# copy our local tempfile to the guest
+
+			puts "DEBUG: Running w/ tools"
+
+			#
+			# Copy our local tempfile to the guest
+			#
 			vmrunstr = "vmrun -T ws -gu #{@vm_user} -gp #{@vm_pass} " +
 					"copyFileFromHostToGuest \'#{@location}\' \'#{local_tempfile_path}\'" +
-					" \'#{remote_tempfile_path}\' nogui"
+					" \'#{remote_tempfile_path}\'"
 			system_command(vmrunstr)
 
-			# now run it on the guest
-			vmrunstr = "vmrun -T ws -gu #{@vm_user} -gp #{@vm_pass} " + 
-					"runProgramInGuest \'#{@location}\' -noWait -activeWindow \'#{remote_run_command}\'"
-			system_command(vmrunstr)
-
-			## CLEANUP
-			# delete it on the guest
+			if @os == "linux"
+				#
+				# Now run the command directly on the guest
+				#
+				vmrunstr = "vmrun -T ws -gu #{@vm_user} -gp #{@vm_pass} " + 
+						"runProgramInGuest \'#{@location}\' /bin/sh #{remote_tempfile_path}"
+				system_command(vmrunstr)
+			else
+				#
+				# Now run the command directly on the guest
+				#
+				vmrunstr = "vmrun -T ws -gu #{@vm_user} -gp #{@vm_pass} " + 
+						"runProgramInGuest \'#{@location}\' #{remote_tempfile_path}"
+				system_command(vmrunstr)
+			end
+			#
+			# Cleanup. Delete it on the guest
+			#
 			vmrunstr = "vmrun -T ws -gu #{@vm_user} -gp #{@vm_pass} " + 
 					"deleteFileInGuest \'#{@location}\' \'#{remote_tempfile_path}\'"
 			system_command(vmrunstr)
 
-			# delete it locally
+			#
+			# Delete it locally
+			#
 			local_delete_command = "rm #{local_tempfile_path}"
 			system_command(local_delete_command)
 		else
-			# since we can't copy easily w/o tools, let's just run it directly :/
+
+			#
+			# Use SCP / SSH
+			#
+
 			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)}"
 
-				output_file = "/tmp/lab_command_output_#{rand(1000000)}"
-
+				#
+				# Copy it over
+				#
 				scp_to(local_tempfile_path, remote_tempfile_path)
-				ssh_exec(command + "> #{output_file}")
-				scp_from(output_file, output_file)
-
-				ssh_exec("rm #{output_file}")
+
+				#
+				# And ... execute it
+				#
+				ssh_exec("/bin/sh #{remote_tempfile_path} > #{remote_output_file}")
+
+				#
+				# Now copy the output back to us
+				#
+				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}")
-
-				# Ghettohack!
-				string = File.open(output_file,"r").read
-				`rm #{output_file}`
-
+
+				# Now, let's look at the output of the command
+				string = File.open(local_output_file,"r").read
+
+				#
+				# And clean that up too
+				#
+				`rm #{local_output_file}`
+
 			else
-				raise "zomgwtfbbqnotools"
-			end	
+				raise "Hey, no tools, and windows? can't do nuttin for ya man."
+			end
+
 		end
 	return string
 	end
@@ -143,7 +195,7 @@ def check_file_exists(file)
 		file = filter_input(file)
 		if @tools
 			vmrunstr = "vmrun -T ws -gu \'#{@vm_user}\' -gp \'#{@vm_pass}\' fileExistsInGuest " +
-				"\'#{@location}\' \'#{file}\' "
+				"\'#{@location}\' \'#{file}\'"
 			system_command(vmrunstr)
 		else
 			raise "Unsupported"
@@ -159,8 +211,6 @@ def create_directory(directory)
 		else
 			raise "Unsupported"
 		end
-
-
 	end
 
 	def cleanup

lib/lab/version.rb

@@ -1,3 +1,3 @@
 module Lab
-  VERSION = "0.1.0"
+  VERSION = "0.1.5"
 end

lib/lab/vm.rb

@@ -49,10 +49,10 @@ def initialize(config = {})
 
 		@location = filter_input(config['location'])
 		#@name = config['name'] || ""
-		@description = config['description'] || ""
-		@tools = config['tools'] || ""
-		@os = config['os'] || ""			
-		@arch = config['arch'] || ""
+		@description = config['description']
+		@tools = config['tools']
+		@os = config['os']
+		@arch = config['arch']
 		@type = filter_input(config['type']) || "unspecified"
 		@credentials = config['credentials'] || []
 
@@ -91,9 +91,9 @@ def initialize(config = {})
 		elsif @driver_type == "remote_workstation"
 			@driver = Lab::Drivers::RemoteWorkstationDriver.new(config)
 		#elsif @driver_type == "qemu"
-		#	@driver = Lab::Drivers::QemuDriver.new	
+		#	@driver = Lab::Drivers::QemuDriver.new
 		#elsif @driver_type == "qemudo"
-		#	@driver = Lab::Drivers::QemudoDriver.new	
+		#	@driver = Lab::Drivers::QemudoDriver.new
 		else
 			raise "Unknown Driver Type"
 		end