#!/usr/bin/env ruby

# -------------------------------------------------------------------------- #
# Copyright 2002-2011, OpenNebula Project Leads (OpenNebula.org)             #
#                                                                            #
# Licensed under the Apache License, Version 2.0 (the "License"); you may    #
# not use this file except in compliance with the License. You may obtain    #
# a copy of the License at                                                   #
#                                                                            #
# http://www.apache.org/licenses/LICENSE-2.0                                 #
#                                                                            #
# Unless required by applicable law or agreed to in writing, software        #
# distributed under the License is distributed on an "AS IS" BASIS,          #
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   #
# See the License for the specific language governing permissions and        #
# limitations under the License.                                             #
#--------------------------------------------------------------------------- #

require 'pp'

require 'rexml/document'
require 'open3'


module HyperV

    STATES=%w{Running Stopped Paused Suspended Starting Snapshotting
        Saving Stopping}

    STATE_TABLE={
        'Running'       => 'a',
        'Stopped'       => '-',
        'Paused'        => 'p',
        'Suspended'     => 'p',
        'Starting'      => 'a',
        'Snapshotting'  => 'a',
        'Saving'        => 'a',
        'Stopping'      => 'a',
        nil             => 'd'
    }

    # This class parses OpenNebula xml templates and generates the commands
    # to create disks and networks
    class Template
        # Initializes the object with the data from template
        #
        # @param [String, #read] template xml definition from deployment.0
        # @param [String] prefix the directory where VM_DIR is located on the
        #   windows machines
        def initialize(template, prefix)
            @prefix=prefix
            @xml=REXML::Document.new(template)
            @root=@xml.root.elements['/TEMPLATE']
            @data=Hash.new
            parse
        end

        # Extracts useful data from the VM definition
        def parse
            @data[:cpu]=get_value('VCPU')
            @data[:memory]=get_value('MEMORY')
            @data[:id]=get_value('VMID')
            @data[:context]=get_value('CONTEXT/TARGET')

            disks=Array.new

            @root.elements.each('DISK') do |d|

                disk=Hash.new
                disk[:id]=get_value('DISK_ID', nil, d)
                disk[:target]=get_value('TARGET', nil, d)

                disks<<disk
            end

            @data[:disks]=disks

            nics=Array.new

            @root.elements.each('NIC') do |n|
                nic=Hash.new
                nic[:bridge]=get_value('BRIDGE', nil, n)
                nic[:mac]=get_value('MAC', nil, n)
                nics<<nic
            end

            @data[:nics]=nics
        end

        # Extracts a value from the xml
        #
        # @param [String] name name of an xml element (XPATH)
        # @param [String] default value it will have if the element does not
        #   exist
        # @param [REXML::Document] root_ use another xml root
        def get_value(name, default=nil, root_=nil)
            root=root_||@root

            value=root.elements["#{name}"]
            if value
                value.text
            else
                default
            end
        end

        # returns the prefix
        #
        # @return [String] VM_DIR directory on the windows machines
        def prefix
            @prefix
        end

        # memory in megabytes
        def memory
            @data[:memory].to_i
        end

        # cpu value from the VM definition
        def cpu
            @data[:cpu].to_i
        end

        def name
            "one-#{@data[:id]}"
        end

        def cmd_create
            "New-VM -Name #{name}"
        end

        def cmd_memory
            "Set-VMMemory #{name} #{memory}MB"
        end

        def cmd_disk(num)
            disk=@data[:disks][num]

            letter=disk[:target][-1,1]
            num=letter.unpack('C')[0]-?a

            controller=num>>1
            disk_num=num&1

            [controller, disk_num]

            "Add-VMDisk #{name} #{controller} #{disk_num} "<<
                "#{prefix}\\#{@data[:id]}\\images\\disk.#{disk[:id]}"
        end

        def cmd_disks
            (0..(@data[:disks].length-1)).map do |n|
                cmd_disk(n)
            end
        end

        def cmd_context
            return nil if !@data[:context]

            controller=1
            disk_num=0

            "Add-VMDisk #{name} #{controller} #{disk_num} "<<
                "#{prefix}\\#{@data[:id]}\\images\\disk.1.iso "<<
                "-DVD"
        end

        def cmd_nic(num)
            nic=@data[:nics][num]

            bridge=nic[:bridge]

            mac=nic[:mac]

            if mac
                mac.tr!(':', '')

                # Swap bytes from each word (tempative workawound to
                # a MAC uncorrectly set)
                #m=mac.split(':')
                #mac=""
                #mac<<m[1]
                #mac<<m[0]
                #mac<<m[3]
                #mac<<m[2]
                #mac<<m[5]
                #mac<<m[4]

                "Add-VMNic #{name} #{bridge} \"#{mac}\" -legacy"
            else
                "Add-VMNic #{name} #{bridge} -legacy"
            end
        end

        def cmd_nics
            (0..(@data[:nics].length-1)).map do |n|
                cmd_nic(n)
            end
        end

        def cmd_start_vm
            "Start-VM #{name} -Wait"
        end
    end

    class Controller
        def initialize(vmid, vmdir, proxy=nil)
            @id=vmid
            @proxy=proxy
            @vmdir=vmdir
        end

        def set_vmid(vmid)
            @id=vmid
        end

        def deploy(host, template)
            vm=Template.new(template, @vmdir)

            ssh(@proxy, hyperv_command(host, vm.cmd_create))
            ssh(@proxy, hyperv_command(host, vm.cmd_memory))

            vm.cmd_disks.each do |disk|
                ssh(@proxy, hyperv_command(host, disk))
            end

            vm.cmd_nics.each do |nic|
                ssh(@proxy, hyperv_command(host, nic))
            end

            context=vm.cmd_context
            ssh(@proxy, hyperv_command(host, context)) if context

            ssh(@proxy, hyperv_command(host, vm.cmd_start_vm))

            vm.name
        end

        def poll(host, identifier)
            text=ssh(host, hyperv_command(host, "Get-VMState #{identifier}"))

            # Strip windowsline ends
            text.gsub!("\r", '')
            text.strip!

            data=Hash[text.scan(/^([^:]+?) *: *(.*)$/)]

            state=STATE_TABLE[data['EnabledState']]
            memory=data['MemoryUsage'].to_i*1024
            cpu_load=data['CPULoad'].to_i/100.0
            cpu_max=data['CPUCount'].to_i*100
            cpu=(cpu_max*cpu_load).to_i

            %{STATE=#{state} USEDMEMORY=#{memory} USEDCPU=#{cpu}}
        end

        def cancel(host, identifier)
            ssh(host, hyperv_command(host, "Stop-VM #{identifier} -Force -Wait"))
            sleep 1
            delete(host, identifier)
        end

        def shutdown(host, identifier)
            cancel(host, identifier)
        end

        def delete(host, identifier)
            ssh(host, hyperv_command(host, "Remove-VM #{identifier} -Force"))
        end

        def monitor(host)
            cpu=host_cpu_info(host)
            memory=host_memory_info(host)

            total_memory=memory[:total_mem]
            free_memory=memory[:free_mem]
            used_memory=total_memory-free_memory

            total_cpu=cpu[:cores]*100
            used_cpu=cpu[:cpu_load]
            free_cpu=total_cpu-used_cpu

            text=<<EOT
HYPERVISOR=hyperv
TOTALCPU=#{total_cpu}
TOTALMEMORY=#{total_memory}
FREEMEMORY=#{free_memory}
USEDMEMORY=#{used_memory}
USEDCPU=#{used_cpu}
FREECPU=#{free_cpu}
EOT
        end

        def host_cpu_info(host)
            text=ssh(host,
                ps("Get-WmiObject win32_processor -ComputerName #{host} | "<<
                "select LoadPercentage,NumberOfCores | fl"))

            text.gsub!("\r", '')

            cores=0
            cpu_load=0
            processors=0
            cores_per_processor=0
            text.split("\n").each do |line|
                m=line.match(/^(\w+)\s*:\s*(\d*)/)

                next if !m

                case m[1]
                when 'LoadPercentage'
                    cpu_load+=m[2].to_i
                when 'NumberOfCores'
                    processors+=1
                    cores_per_processor=m[2].to_i
                    cores+=cores_per_processor
                end
            end

           cpu_load=(cpu_load/(processors*100.0))*(cores*100)

            {
                :cores => cores,
                :processors => processors,
                :cpu_load => cpu_load,
                :cores_per_processor => cores_per_processor
            }
        end

        def host_memory_info(host)
            total_text=ssh(host,
                ps("(Get-WmiObject -Class Win32_ComputerSystem "<<
                "-ComputerName #{host}).TotalPhysicalMemory"))
            free_text=ssh(host,
                ps("(Get-WmiObject -Class Win32_OperatingSystem "<<
                "-ComputerName #{host}).FreePhysicalMemory"))
            {
                :total_mem => total_text.strip.to_i/1024,
                :free_mem => free_text.strip.to_i
            }
        end

        def hyperv_command(host, command)
             ps("#{command} -server #{host}")
        end

        def ps(command)
            cmd=command.gsub("\\", "\\\\\\\\\\")
            "powershell -command \"\\\\\\\$error.clear() ; #{cmd} ; " <<
                "if (\\\\\\\$error.count -ne 0) " <<
                "{ Write-Host 'ONE-ERROR-CODE: Error' }\""
        end

        def ssh(host_, command)
            host=@proxy||host_
            sanitized_command=command.gsub!('"', "\\\"")
            cmd="ssh -n #{host} \"#{sanitized_command}\""
            STDERR.puts cmd
            stdin, stdout, stderr=Open3.popen3(cmd)

            stdin.close
            out=stdout.read
            err=stderr.read

            if out.match(/^ONE-ERROR-CODE: Error/)
                STDERR.puts "==== STDOUT ===="
                STDERR.puts out
                STDERR.puts "==== STDERR ===="
                STDERR.puts err

                exit(-1)
            else
                out
            end
        end

    end
end

vm_dir=ENV['ONE_HYPERV_VMDIR']
proxy=ENV['ONE_HYPERV_PROXY']

if !vm_dir
    STDERR.puts "ONE_HYPERV_VMDIR not set"
    exit(-1)
end

command, host, identifier=ARGV[0..2]
controller=HyperV::Controller.new(nil, vm_dir, proxy)

case command
when 'deploy'
    puts controller.deploy(host, File.read(identifier))
when 'poll'
    puts controller.poll(host, identifier)
when 'cancel'
    puts controller.cancel(host, identifier)
when 'monitor'
    puts controller.monitor(host)
else
    puts "Command not available"
end


