#!/usr/bin/ruby

require 'rbvmomi'
require 'redis'
require 'time'
require 'yaml'

$:.unshift(File.dirname(__FILE__))
require 'lib/logger'
require 'lib/require_relative'
require 'lib/vsphere_helper'

Dir.chdir(File.dirname(__FILE__))

# Load the configuration file
config_file = File.expand_path('vmware-host-pooler.yaml')
config = YAML.load_file(config_file)

pools   = config[:pools]
vsphere = config[:vsphere]

# Load logger library
$logger = Logger.new config[:config]['logfile']

# Load Graphite helper library (if configured)
if (defined? config[:config]['graphite'])
  require 'lib/graphite'
  $graphite = Graphite.new config[:config]['graphite']
end

# Connect to Redis
$redis = Redis.new

# vSphere object
$vsphere = {}

# Our thread-tracker object
$threads = {}



# Check the state of a VM
def check_pending_vm vm, pool, timeout
  Thread.new {
    host = $vsphere[pool].find_vms(vm)[vm]

    if (host)
      if (
        (host.summary) and
        (host.summary.guest) and
        (host.summary.guest.hostName) and
        (host.summary.guest.hostName == vm)
      )
        begin
          Socket.getaddrinfo(vm, nil)
        rescue
        end

        $redis.smove('vmware_host_pool__pending__'+pool, 'vmware_host_pool__ready__'+pool, vm)

        $logger.log('s', "[>] [#{pool}] '#{vm}' moved to 'ready' queue")
      else
        if (
          (host.runtime) and
          (host.runtime.bootTime) and
          ((((Time.now - host.runtime.bootTime)/60).to_s[/^\d+\.\d{1}/].to_f) > timeout)
        )
          $redis.smove('vmware_host_pool__pending__'+pool, 'vmware_host_pool__completed__'+pool, vm)

          $logger.log('d', "[!] [#{pool}] '#{vm}' marked as 'failed' after #{timeout} minutes")
        end
      end
    end
  }
end

def check_ready_vm vm, pool, ttl
  Thread.new {
    host = $vsphere[pool].find_vms(vm)[vm]

    if (host)
      if (
        (host.runtime) and
        (host.runtime.powerState) and
        (host.runtime.powerState != 'poweredOn')
      )
        $redis.smove('vmware_host_pool__ready__'+pool, 'vmware_host_pool__completed__'+pool, vm)

        $logger.log('d', "[!] [#{pool}] '#{vm}' appears to be powered off or dead")
      else
        if ((((Time.now - host.runtime.bootTime)/60).to_s[/^\d+\.\d{1}/].to_f) > ttl)
          $redis.smove('vmware_host_pool__ready__'+pool, 'vmware_host_pool__completed__'+pool, vm)

          $logger.log('d', "[!] [#{pool}] '#{vm}' reached end of TTL after #{ttl} minutes")
        end
      end
    end
  }
end

def check_running_vm vm, pool, ttl
  Thread.new {
    host = $vsphere[pool].find_vms(vm)[vm]

    if (host)
      if (
        (host.runtime) and
        (host.runtime.powerState != 'poweredOn')
      )
        $redis.smove('vmware_host_pool__running__'+pool, 'vmware_host_pool__completed__'+pool, vm)

        $logger.log('d', "[!] [#{pool}] '#{vm}' appears to be powered off or dead")
      else
        if (
          (host.runtime) and
          (host.runtime.bootTime)
          ((((Time.now - host.runtime.bootTime)/60).to_s[/^\d+\.\d{1}/].to_f) > ttl)
        )
          $redis.smove('vmware_host_pool__running__'+pool, 'vmware_host_pool__completed__'+pool, vm)

          $logger.log('d', "[!] [#{pool}] '#{vm}' reached end of TTL after #{ttl} minutes")
        end
      end
    end
  }
end

# Clone a VM
def clone_vm template, pool, folder, datastore
  Thread.new {
    vm = {}

    if template =~ /\//
      templatefolders = template.split('/')
      vm['template'] = templatefolders.pop
    end

    if templatefolders
      vm[vm['template']] = $vsphere[vm['template']].find_folder(templatefolders.join('/')).find(vm['template'])
    else
      raise "Please provide a full path to the template"
    end

    if vm['template'].length == 0
      raise "Unable to find template '#{vm['template']}'!"
    end

    # Generate a randomized hostname
    o = [('a'..'z'),('0'..'9')].map{|r| r.to_a}.flatten
    vm['hostname'] = o[rand(25)]+(0...14).map{o[rand(o.length)]}.join

    # Add VM to Redis inventory ('pending' pool)
    $redis.sadd('vmware_host_pool__pending__'+vm['template'], vm['hostname'])

    # Annotate with creation time, origin template, etc.
    configSpec = RbVmomi::VIM.VirtualMachineConfigSpec(
      :annotation =>
        'Base template: ' + vm['template'] + "\n" +
        'Creation time: ' + Time.now.strftime("%Y-%m-%d %H:%M")
    )

    # Put the VM in the specified folder and resource pool
    relocateSpec = RbVmomi::VIM.VirtualMachineRelocateSpec(
      :datastore    => $vsphere[vm['template']].find_datastore(datastore),
      :pool         => $vsphere[vm['template']].find_pool(pool),
      :diskMoveType => :moveChildMostDiskBacking
    )

    # Create a clone spec
    spec = RbVmomi::VIM.VirtualMachineCloneSpec(
      :location      => relocateSpec,
      :config        => configSpec,
      :powerOn       => true,
      :template      => false
    )

    # Clone the VM
    $logger.log('d', "[ ] [#{vm['template']}] '#{vm['hostname']}' is being cloned from '#{vm['template']}'")

    start = Time.now
    vm[vm['template']].CloneVM_Task(
      :folder => $vsphere[vm['template']].find_folder(folder),
      :name => vm['hostname'],
      :spec => spec
    ).wait_for_completion
    finish = '%.2f' % (Time.now-start)

    $logger.log('s', "[+] [#{vm['template']}] '#{vm['hostname']}' cloned from '#{vm['template']}' in #{finish} seconds")

    begin
      $graphite.log("vcloud.clone.#{vm['template']}", finish) if defined? $graphite
    rescue
    end
  }
end

# Destroy a VM
def destroy_vm vm, pool
  Thread.new {
    $redis.srem('vmware_host_pool__completed__'+pool, vm)
    $redis.hdel('vmware_host_pool__active__'+pool, vm)

    host = $vsphere[pool].find_vms(vm)[vm]

    if (host)
      start = Time.now

      if (
        (host.runtime) and
        (host.runtime.powerState) and
        (host.runtime.powerState == 'poweredOn')
      )
        $logger.log('d', "[ ] [#{pool}] '#{vm}' is being shut down")
        host.PowerOffVM_Task.wait_for_completion
      end

      host.Destroy_Task.wait_for_completion
      finish = '%.2f' % (Time.now-start)

      $logger.log('s', "[-] [#{pool}] '#{vm}' destroyed in #{finish} seconds")

      $graphite.log("vcloud.destroy.#{pool}", finish) if defined? $graphite
    end
  }
end

def check_pool pool
  $logger.log('d', "[*] [#{pool['name']}] starting worker thread")

  $redis.del('vmware_host_pool__pending__'+pool['name'])

  $threads[pool['name']] = Thread.new {
    $vsphere[pool['name']] ||= VsphereHelper.new

    loop do
      # INVENTORY
      inventory = {}
      begin
        base = $vsphere[pool['name']].find_pool(pool['pool'])

        base.vm.each do |vm|
          if (
            (! $redis.sismember('vmware_host_pool__running__'+pool['name'], vm['name'])) and
            (! $redis.sismember('vmware_host_pool__ready__'+pool['name'], vm['name'])) and
            (! $redis.sismember('vmware_host_pool__pending__'+pool['name'], vm['name'])) and
            (! $redis.sismember('vmware_host_pool__completed__'+pool['name'], vm['name'])) and
            (! $redis.sismember('vmware_host_pool__discovered__'+pool['name'], vm['name']))
          )
            $redis.sadd('vmware_host_pool__discovered__'+pool['name'], vm['name'])

            $logger.log('s', "[?] [#{pool['name']}] '#{vm['name']}' added to 'discovered' queue")
          end

          inventory[vm['name']] = 1
        end
      rescue
      end

      # RUNNING
      $redis.smembers('vmware_host_pool__running__'+pool['name']).each do |vm|
        if (inventory[vm])
          if (pool['running_ttl'])
            begin
              check_running_vm(vm, pool['name'], pool['running_ttl'])
            rescue
            end
          else
            begin
              check_running_vm(vm, pool['name'], '720')
            rescue
            end
          end
        end
      end

      # READY
      $redis.smembers('vmware_host_pool__ready__'+pool['name']).each do |vm|
        if (inventory[vm])
          if (pool['ready_ttl'])
            begin
              check_ready_vm(vm, pool['name'], pool['ready_ttl'])
            rescue
            end
          end
        end
      end

      # PENDING
      $redis.smembers('vmware_host_pool__pending__'+pool['name']).each do |vm|
        pool['timeout'] ||= 10

        if (inventory[vm])
          begin
            check_pending_vm(vm, pool['name'], pool['timeout'])
          rescue
          end
        end
      end

      # COMPLETED
      $redis.smembers('vmware_host_pool__completed__'+pool['name']).each do |vm|
        if (inventory[vm])
          begin
            destroy_vm(vm, pool['name'])
          rescue
          end
        end
      end

      # LONG-RUNNING
      $redis.smembers('vmware_host_pool__running__'+pool['name']).each do |vm|
        if ($redis.hget('vmware_host_pool__active__'+pool['name'], vm))
          running = (Time.now - Time.parse($redis.hget('vmware_host_pool__active__'+pool['name'], vm)))/60/60
          if (running > 12)
            $redis.smove('vmware_host_pool__running__'+pool['name'], 'vmware_host_pool__completed__'+pool['name'], vm)

            $logger.log('d', "[!] [#{pool['name']}] '#{vm}' reached end of TTL after 12 hours")
          end
        end
      end

      # REPOPULATE
      total = $redis.scard('vmware_host_pool__ready__'+pool['name']) +
              $redis.scard('vmware_host_pool__pending__'+pool['name'])

      begin
        $graphite.log(
          'vcloud.ready.'+pool['name'], $redis.scard('vmware_host_pool__ready__'+pool['name'])
        ) if defined? $graphite
      rescue
      end

      if (total < pool['size'])
        (1..(pool['size'] - total)).each { |i|
          begin
            clone_vm(
              pool['template'],
              pool['pool'],
              pool['folder'],
              pool['datastore']
            )
          rescue
          end
        }
      end

      sleep(1)
    end
  }
end



$logger.log('d', "starting vmware-host-pooler")

loop do
  pools.each do |pool|
    if (! $threads[pool['name']])
      check_pool(pool)
    else
      if (! $threads[pool['name']].alive?)
        $logger.log('d', "[!] [#{pool['name']}] worker thread died, restarting")
        check_pool(pool)
      end
    end
  end

  sleep(1)
end