vm_guest.rb in rvc-1.8.0

- old
+ new

@@ -1,33 +1,57 @@
+# Copyright (c) 2013 VMware, Inc.  All Rights Reserved.
+# 
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# 
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+# 
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+require 'rvc/vim'
+
 opts :authenticate do
   summary "Authenticate within guest"
-  arg :vm, nil, :lookup => VIM::VirtualMachine
+  arg :vm, nil, :lookup => VIM::VirtualMachine, :multi => true
   opt :interactive_session, "Allow command to interact with desktop", :default => false, :type => :bool
   opt :password, "Password in guest", :type => :string
   opt :username, "Username in guest", :default => "root", :type => :string
 end
 
-def authenticate vm, opts
-  auth = ((@auths ||= {})[vm] ||= {})[opts[:username]]
-
-  if opts[:password].nil? or opts[:password].empty?
-    opts[:password] = ask("password: ") { |q| q.echo = false }
+def authenticate vms, opts
+  vms.each do |vm|
+    auth = ((@auths ||= {})[vm] ||= {})[opts[:username]]
+  
+    if opts[:password].nil? or opts[:password].empty?
+      opts[:password] = ask("password: ") { |q| q.echo = false }
+    end
+  
+    auth = VIM.NamePasswordAuthentication(
+      :username => opts[:username],
+      :password => opts[:password],
+      :interactiveSession => opts[:interactive_session]
+    )
+  
+    @auths[vm][opts[:username]] = auth
+    begin
+      check_auth vm, opts
+    rescue
+      clear_auth vm, opts
+      err "Could not authenticate: #{$!}"
+    end
   end
-
-  auth = VIM.NamePasswordAuthentication(
-    :username => opts[:username],
-    :password => opts[:password],
-    :interactiveSession => opts[:interactive_session]
-  )
-
-  @auths[vm][opts[:username]] = auth
-  begin
-    check_auth vm, opts
-  rescue
-    clear_auth vm, opts
-    err "Could not authenticate: #{$!}"
-  end
 end
 
 
 opts :check_auth do
   summary "Check credentials"
@@ -238,11 +262,72 @@
       :auth => auth,
       :filePath => opts[:guest_path]
     )
 end
 
+def generic_http_download uri, local_path
+  http = Net::HTTP.new(uri.host, uri.port)
+  http.use_ssl = true
+  http.verify_mode = OpenSSL::SSL::VERIFY_NONE
+  #http.set_debug_output $stderr
+  http.start
 
+#  headers = { 'cookie' => connection.cookie }
+  headers = {}
+  http_path = "#{uri.path}?#{uri.query}"
+  http.request_get(http_path, headers) do |res|
+    case res
+    when Net::HTTPOK
+      len = res.content_length
+      count = 0
+      File.open(local_path, 'wb') do |io|
+        res.read_body do |segment|
+          count += segment.length
+          io.write segment
+          $stdout.write "\e[0G\e[Kdownloading #{count}/#{len} bytes (#{(count*100)/len}%)"
+          $stdout.flush
+        end
+      end
+      $stdout.puts
+    else
+      err "download failed: #{res.message}"
+    end
+  end
+end
+  
+def generic_http_upload local_path, uri, size = nil
+  http = Net::HTTP.new(uri.host, uri.port)
+  http.use_ssl = true
+  http.verify_mode = OpenSSL::SSL::VERIFY_NONE
+  #http.set_debug_output $stderr
+  http.start
+
+
+  open(local_path, 'rb') do |io|
+    stream = ProgressStream.new(io, size || io.stat.size) do |s|
+      $stdout.write "\e[0G\e[Kuploading #{s.count}/#{s.len} bytes (#{(s.count*100)/s.len}%)"
+      $stdout.flush
+    end
+
+    headers = {
+      'content-length' => (size || io.stat.size).to_s,
+      'Content-Type' => 'application/octet-stream',
+    }
+    http_path = "#{uri.path}?#{uri.query}"
+
+    request = Net::HTTP::Put.new http_path, headers
+    request.body_stream = stream
+    res = http.request(request)
+    $stdout.puts
+    case res
+    when Net::HTTPOK
+    else
+      err "upload failed: #{res.message}"
+    end
+  end
+end
+
 opts :download_file do
   summary "Download file from guest"
   arg :vm, nil, :lookup => VIM::VirtualMachine
   opt :guest_path, "Path in guest to download from", :required => true, :type => :string
   opt :local_path, "Local file to download to", :required => true, :type => :string
@@ -263,56 +348,66 @@
       :guestFilePath => opts[:guest_path]
     ).url
 
   download_uri = URI.parse(download_url.gsub /http(s?):\/\/\*:[0-9]*/, "")
   download_path = "#{download_uri.path}?#{download_uri.query}"
-
-  http_download vm._connection, download_path, opts[:local_path]
+  
+  generic_http_download download_uri, opts[:local_path]
 end
 
 
 opts :upload_file do
   summary "Upload file to guest"
-  arg :vm, nil, :lookup => VIM::VirtualMachine
+  arg :vm, nil, :lookup => VIM::VirtualMachine, :multi => true
   opt :group_id, "Group ID of file", :type => :int
   opt :guest_path, "Path in guest to upload to", :required => true, :type => :string
   opt :local_path, "Local file to upload", :required => true, :type => :string
   opt :overwrite, "Overwrite file", :default => false, :type => :bool
   opt :owner_id, "Owner ID of file", :type => :int
   opt :permissions, "Permissions of file", :type => :string
   opt :username, "Username in guest", :default => "root", :type => :string
 end
 
-def upload_file vm, opts
-  guestOperationsManager = vm._connection.serviceContent.guestOperationsManager
-  err "This command requires vSphere 5 or greater" unless guestOperationsManager.respond_to? :fileManager
-  fileManager = guestOperationsManager.fileManager
-
-  opts[:permissions] = opts[:permissions].to_i(8) if opts[:permissions]
-
-  auth = get_auth vm, opts
-
-  file = File.new(opts[:local_path], 'rb')
-
-  upload_url = fileManager.
-    InitiateFileTransferToGuest(
-      :vm => vm,
-      :auth => auth,
-      :guestFilePath => opts[:guest_path],
-      :fileAttributes => VIM.GuestPosixFileAttributes(
-        :groupId => opts[:group_id],
-        :ownerId => opts[:owner_id],
-        :permissions => opts[:permissions]
-      ),
-      :fileSize => file.size,
-      :overwrite => opts[:overwrite]
-    )
-
-  upload_uri = URI.parse(upload_url.gsub /http(s?):\/\/\*:[0-9]*/, "")
-  upload_path = "#{upload_uri.path}?#{upload_uri.query}"
-
-  http_upload vm._connection, opts[:local_path], upload_path
+def upload_file vms, opts
+  vms.each do |vm|
+    guestOperationsManager = vm._connection.serviceContent.guestOperationsManager
+    err "This command requires vSphere 5 or greater" unless guestOperationsManager.respond_to? :fileManager
+    fileManager = guestOperationsManager.fileManager
+  
+    opts[:permissions] = opts[:permissions].to_i(8) if opts[:permissions]
+  
+    auth = get_auth vm, opts
+  
+    file_size = nil
+    if opts[:local_path] =~ /^http:/
+      # XXX: Not acceptable for big files
+      file_size = open(opts[:local_path], 'rb').read.length
+    else
+      err "local file does not exist" unless File.exists? local_path
+      file = File.new(opts[:local_path], 'rb')
+      file_size = file.size
+    end  
+  
+    upload_url = fileManager.
+      InitiateFileTransferToGuest(
+        :vm => vm,
+        :auth => auth,
+        :guestFilePath => opts[:guest_path],
+        :fileAttributes => VIM.GuestPosixFileAttributes(
+          :groupId => opts[:group_id],
+          :ownerId => opts[:owner_id],
+          :permissions => opts[:permissions]
+        ),
+        :fileSize => file_size,
+        :overwrite => opts[:overwrite]
+      )
+  
+    upload_uri = URI.parse(upload_url.gsub /http(s?):\/\/\*:[0-9]*/, "")
+    upload_path = "#{upload_uri.path}?#{upload_uri.query}"
+  
+    generic_http_upload opts[:local_path], upload_uri, file_size
+  end
 end
 
 
 opts :ls_guest do
   summary "List files in guest"
@@ -429,11 +524,11 @@
 
 
 # Process commands
 opts :start_program do
   summary "Run program in guest"
-  arg :vm, nil, :lookup => VIM::VirtualMachine
+  arg :vm, nil, :lookup => VIM::VirtualMachine, :multi => true
   opt :arguments, "Arguments of command", :default => "", :type => :string
   opt :background, "Don't wait for process to finish", :default => false, :type => :bool
   opt :delay, "Interval in seconds", :type => :float, :default => 5.0
   opt :env, "Environment variable(s) to set (e.g. VAR=value)", :multi => true, :type => :string
   opt :program_path, "Path to program in guest", :required => true, :type => :string
@@ -442,49 +537,53 @@
   opt :working_directory, "Working directory of the program to run", :type => :string
   conflicts :background, :timeout
   conflicts :background, :delay
 end
 
-def start_program vm, opts
-  guestOperationsManager = vm._connection.serviceContent.guestOperationsManager
-  err "This command requires vSphere 5 or greater" unless guestOperationsManager.respond_to? :processManager
-  processManager = guestOperationsManager.processManager
-
-  auth = get_auth vm, opts
-
-  pid = processManager.
-    StartProgramInGuest(
-      :vm => vm,
-      :auth => auth,
-      :spec => VIM.GuestProgramSpec(
-        :arguments => opts[:arguments],
-        :programPath => opts[:program_path],
-        :envVariables => opts[:env],
-        :workingDirectory => opts[:working_directory]
-      )
-    )
-
-  Timeout.timeout opts[:timeout] do
-    while true
-      processes = processManager.
-        ListProcessesInGuest(
-          :vm => vm,
-          :auth => auth,
-          :pids => [pid]
+def start_program vms, opts
+  vms.each do |vm|
+    guestOperationsManager = vm._connection.serviceContent.guestOperationsManager
+    err "This command requires vSphere 5 or greater" unless guestOperationsManager.respond_to? :processManager
+    processManager = guestOperationsManager.processManager
+  
+    auth = get_auth vm, opts
+  
+    pid = processManager.
+      StartProgramInGuest(
+        :vm => vm,
+        :auth => auth,
+        :spec => VIM.GuestProgramSpec(
+          :arguments => opts[:arguments],
+          :programPath => opts[:program_path],
+          :envVariables => opts[:env],
+          :workingDirectory => opts[:working_directory]
         )
-      process = processes.first
-
-      if !process.endTime.nil?
-        if process.exitCode != 0
-          err "Process failed with exit code #{process.exitCode}"
+      )
+  
+    begin
+      Timeout.timeout opts[:timeout] do
+        while true
+          processes = processManager.
+            ListProcessesInGuest(
+              :vm => vm,
+              :auth => auth,
+              :pids => [pid]
+            )
+          process = processes.first
+    
+          if !process.endTime.nil?
+            if process.exitCode != 0
+              err "Process failed with exit code #{process.exitCode}"
+            end
+            break
+          elsif opts[:background]
+            break
+          end
+    
+          sleep opts[:delay]
         end
-        break
-      elsif opts[:background]
-        break
       end
-
-      sleep opts[:delay]
+    rescue Timeout::Error
+      err "Timed out waiting for process to finish."
     end
   end
-rescue Timeout::Error
-  err "Timed out waiting for process to finish."
 end