#!/usr/bin/env ruby
require 'rubygems'
require 'esx'
require 'terminal-table/import'
require 'clamp'
require 'net/http'
require 'fileutils'

class BaseCommand < Clamp::Command
  parameter "ADDRESS", "ESX host address"
  option "--user", "USER", "Username", :default => "root"
  option "--password", "PASSWORD", "Password", :default => ""
  option "--debug", :flag, "Print debugging info"
end

class InfoCommand < BaseCommand

  parameter "ADDRESS", "ESX host address"
  option "--user", "USER", "Username", :default => "root"
  option "--password", "PASSWORD", "Password", :default => ""

  def execute
    begin

      host = ESX::Host.connect(address, user, password)

      puts
      name = host.name.upcase
      puts "*" * name.size
      puts name
      puts "*" * name.size
      puts "Memory Size:      %s GB" % host.memory_size.bytes.to.gigabytes.to_f.round
      puts "Memory Usage:     %s GB" % host.memory_usage.bytes.to.gigabytes.to_f.round
      puts "Cpu Cores:        %s" % host.cpu_cores
      puts "Power State:      %s" % host.power_state
      puts "Hosted VMs:       %s" % host.virtual_machines.size
      puts "Running VMs:      %s" % (host.virtual_machines.find_all{ |vm| vm.power_state == 'poweredOn' }).size

      if not host.virtual_machines.empty?
        puts "\nVirtual Machines:"
        user_table = table do |t|
          t.headings = "NAME","MEMORY","CPUS","NICS","DISKS", "STATE"
          host.virtual_machines.each do |vm|
            t << [vm.name,vm.memory_size.bytes.to.megabytes.to_i, vm.cpus, vm.ethernet_cards_number, vm.virtual_disks_number, vm.power_state]
          end
        end
        puts user_table
      end

      puts "\nDatastores:"
      user_table = table do |t|
          t.headings = "NAME","CAPACITY","FREESPACE","ACCESIBLE","TYPE","URL"
        host.datastores.each do |ds|
          dsname = ds.name
          if dsname.size > 20
            dsname = dsname[0..19] + '...'
          end
          t << [dsname,ds.capacity, ds.free_space, ds.datastore_type, ds.accessible, ds.url]
        end
      end
      puts user_table
      puts

    rescue Exception => e
      puts "Error connecting to the ESX host"
      puts "\n#{e.message}"
      if debug?
        puts e.backtrace
      end
    end
  end
end

class DestroyVMCommand < BaseCommand

  parameter "ADDRESS", "ESX host address"
  option "--vm-name", "VM_NAME", "Virtual Machine to destroy"

  def execute
    begin
      host = ESX::Host.connect address, user, password
      host.virtual_machines.each do |vm|
        if vm.name == vm_name 
          print "Destroying Virtual Machine '#{vm.name}'..."
          vm.power_off unless vm.power_state == 'poweredOff'
          vm.destroy
          puts "Done."
        end
      end
    rescue Exception => e
      $stderr.puts "Can't connect to the host #{address}."
      if debug?
        $stderr.puts e.message
      end
      exit 1
    end
  end

end

class PowerOffVMCommand < BaseCommand

  parameter "ADDRESS", "ESX host address"
  option "--vm-name", "VM_NAME", "Virtual Machine to Power Off"

  def execute
    begin
      host = ESX::Host.connect address, user, password
      host.virtual_machines.each do |vm|
        if vm.name == vm_name and vm.power_state == 'poweredOn'
          print "Powering Off VM #{vm_name}... "
          vm.power_off
          puts "Done."
        end
      end
    rescue Exception => e
      $stderr.puts "Can't connect to the host #{address}."
      if debug?
        $stderr.puts e.message
      end
      exit 1
    end
  end

end

class PowerOnVMCommand < BaseCommand

  parameter "ADDRESS", "ESX host address"
  option "--vm-name", "VM_NAME", "Virtual Machine to Power On"

  def execute
    begin
      host = ESX::Host.connect address, user, password
      host.virtual_machines.each do |vm|
        if vm.name == vm_name and vm.power_state != 'poweredOn'
          print "Powering On VM #{vm_name}... "
          vm.power_on
          puts "Done."
        end
      end
    rescue Exception => e
      $stderr.puts "Can't connect to the host #{address}."
      if debug?
        $stderr.puts e.message
      end
      exit 1
    end
  end

end


class CreateVMCommand < BaseCommand

  parameter "ADDRESS", "ESX host address"
  option "--disk-file", "DISK_FILE", "VMDK file to import", :attribute_name => :disk_file
  option "--disk-size", "DISK_SIZE", "VM Disk size in MB", :attribute_name => :disk_size, :default => 8192
  option "--guest-id", "GUEST_ID", "GuestID value", :attribute_name => :guest_id, :default => 'otherGuest'
  option "--name", "VM NAME", "Virtual Machine name (required)"
  option "--memory", "MEMORY", "VM Memory size in MB", :default => 512
  option "--mac-address", "MAC", "VM Nic1 MAC address", :default => nil
  option "--vm-network", "VM NETWORK", "Network where nic is attached to", :default => 'VM Network'
  option "--tmpdir", "TMP DIR", "tmp dir used to download files", :default => "/tmp"
  option "--datastore", "DATASTORE", "Datastore used to host the disk", :default => "datastore1"
  option "--poweron", :flag, "Power on the VM after creation"

  def execute
    begin
      host = ESX::Host.connect address, user, password
    rescue Exception => e
      $stderr.puts "Can't connect to the host #{address}."
      if debug?
        $stderr.puts e.message
      end
      exit 1
    end
    if disk_size and disk_file
      $stderr.puts "Both --disk-file and --disk-size specified. --disk-size will be ignored."
    end
    downloaded_file = nil
    if disk_file.nil?
      # if --disk-file nil? create the VM without disk
      vm = host.create_vm :vm_name => name, 
                          :datastore => datastore, :disk_type => :flat, :memory => memory,
                          :disk_size => disk_size,
                          :guest_id => guest_id, :nics => [{:mac_address => mac_address, :network => vm_network}]
    else
      df = disk_file.dup
      if df.strip.chomp =~ /^http/
        begin
          downloaded_file = disk_file.dup
          tmpfile = "#{tmpdir}/#{Time.now.to_i}.esx"
          puts "Downloading file... (#{tmpfile})"
          download! downloaded_file, tmpfile
          puts
          df = tmpfile
        rescue Exception => e
          FileUtils.rm_f(tmpfile) 
          $stderr.puts "Error downloading file from #{downloaded_file}."
          $stderr.puts e.message if debug?
          exit 1
        end
      end
      raise Exception.new("Invalid disk file") if not File.exist?(df)
      if not name
        $stderr.puts "Invalid VM name."
        $stderr.puts "Use --name option to specify the VM name"
        exit 1
      end
      host.remote_command "mkdir /vmfs/volumes/#{datastore}/#{name}"

      begin
        host.import_disk df, "/vmfs/volumes/#{datastore}/#{name}/#{name}.vmdk"
      rescue Exception => e
        $stderr.puts "Error uploading file to /vmfs/volumes/#{datastore}/#{name}/#{name}.vmdk"
        $stderr.puts e.message if debug?
        exit 1
      end

      if not downloaded_file.nil?
        puts "Deleting tmp file #{df}" if debug?
        FileUtils.rm_f(df) 
      end
      vm = host.create_vm :vm_name => name, 
                          :disk_file => "#{name}/#{name}.vmdk", 
                          :datastore => datastore, :disk_type => :flat, :memory => memory,
                          :guest_id => guest_id, :mac_address => mac_address
    end
    if poweron?
      vm.power_on
    end
  end

  def report_progress(progress, total, show_parts=true)
    line_reset = "\r\e[0K" 
    percent = (progress.to_f / total.to_f) * 100
    line = "Progress: #{percent.to_i}%"
    line << " (#{progress} / #{total})" if show_parts
    line = "#{line_reset}#{line}"
    $stdout.sync = true
    $stdout.print line
  end

  def download!(source_url, destination_file)
    dst = File.open(destination_file, 'w')
    proxy_uri = URI.parse(ENV["http_proxy"] || "")
    uri = URI.parse(source_url)
    http = Net::HTTP.new(uri.host, uri.port, proxy_uri.host, proxy_uri.port, proxy_uri.user, proxy_uri.password)

    if uri.scheme == "https"
      http.use_ssl = true
      http.verify_mode = OpenSSL::SSL::VERIFY_NONE
    end

    res = http.start do |h|
      h.request_get(uri.request_uri) do |response|
        total = response.content_length
        progress = 0
        segment_count = 0

        response.read_body do |segment|
          # Report the progress out
          progress += segment.length
          segment_count += 1

          # Progress reporting is limited to every 25 segments just so
          # we're not constantly updating
          if segment_count % 25 == 0
            report_progress(progress, total)
            segment_count = 0
          end
          # Store the segment
          dst.write(segment)
        end
      end
    end

    raise Exception.new("HTTP Error #{res}") if res.class != Net::HTTPOK
  ensure
    dst.close
  end

end

class DefaultCommand < Clamp::Command
  default_subcommand "info", "Display host info", InfoCommand
  subcommand "create-vm", "Create a VM", CreateVMCommand
  subcommand "poweron-vm", "Power On a VM", PowerOnVMCommand
  subcommand "poweroff-vm", "Power Off a VM", PowerOffVMCommand
  subcommand "destroy-vm", "Destroy VM", DestroyVMCommand
end

begin
  DefaultCommand.run
rescue Exception => e
  puts e.message
  if $DEBUG
    puts $!
    puts $@
  end
end