#!/usr/bin/env ruby

# -------------------------------------------------------------------------- #
# Copyright 2002-2012, 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'

require 'rubygems'
require 'winrm'


if ENV['ONE_LOCATION']
    ONE_LIB_PATH=ENV['ONE_LOCATION']+"/lib/ruby"
    ONE_ETC_PATH=ENV['ONE_LOCATION']+"/etc"
else
    ONE_LIB_PATH='/usr/lib/one/ruby'
    ONE_ETC_PATH='/etc/one'
end

$: << ONE_LIB_PATH

require 'scripts_common'


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, vmdir=nil)
            @prefix=prefix
            @vmdir=vmdir
            @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]

            if disk[:target][0,1]=='s'
                controller_size=2
                scsi=' -SCSI'
            else
                controller_size=64
                scsi=''
            end

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

            controller=num/controller_size
            disk_num=(num%controller_size)

            [controller, disk_num]

            "Add-VMDisk #{name} #{controller} #{disk_num} "<<
                hdisk_path(disk[:id])<<scsi
        end

        def shared_hdisk_path(num)
            "#{prefix}\\#{@data[:id]}\\images\\disk.#{num}"
        end

        def local_hdisk_path(num)
            "#{@vmdir}\\#{@data[:id]}\\images\\disk.#{num}"
        end

        def cmd_scsi_controller
            "Add-VMSCSIController #{name}"
        end

        def cmd_copy_disks
            if @vmdir
                array=["mkdir #{@vmdir}\\#{@data[:id]}\\images"]
                array+=["copy "<<
                    shared_hdisk_path(0).gsub(/\\[^\\]+$/, '')<<"\\* "<<
                    local_hdisk_path(0).gsub(/\\[^\\]+$/, '')]
            else
                nil
            end
        end

        def cmd_del_disks
            if @vmdir
                "del -r #{@vmdir}\\#{@data[:id]}\\images"
            else
                nil
            end
        end

        def hdisk_path(num)
            if @vmdir
                local_hdisk_path(num)
            else
                shared_hdisk_path(num)
            end
        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} "<<
                cdisk_path <<
                "-DVD"
        end

        def shared_cdisk_path
            "#{prefix}\\#{@data[:id]}\\images\\disk.1.iso "
        end

        def local_cdisk_path
            "#{@vmdir}\\#{@data[:id]}\\images\\disk.1.iso "
        end

        def cdisk_path
            if @vmdir
                shared_cdisk_path
            else
                local_cdisk_path
            end
        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, local_vmdir=nil)
            @id=vmid
            @proxy=proxy
            @vmdir=vmdir
            @local_vmdir=local_vmdir
        end

        def winrm(host)
            if !@winrm
                h=@proxy||host
                endpoint="http://#{host}:5985/wsman"
                @winrm=WinRM::WinRMWebService.new(endpoint, :plaintext, :user => USER, :pass=>PASSWORD, :basic_auth_only => true)
             end

             @winrm
        end

        def set_vmid(vmid)
            @id=vmid
        end

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

            response=hyperv_command(host, vm.cmd_create)
            error_if(response, "Error creating VM")

            response=hyperv_command(host, vm.cmd_memory)
            error_if(response, "Error setting VM memory")

            copy=vm.cmd_copy_disks
            if copy
                STDERR.puts "Copying disks"
                copy.each do |line|
                    STDERR.puts "command: #{line}"
                    response=ps(host, line)
                    error_if(response, "Error copying VM images")
                end
            end

            # Create SCSI controller
            hyperv_command(host, vm.cmd_scsi_controller)

            vm.cmd_disks.each do |disk|
                response=hyperv_command(host, disk)
                error_if(response, "Error adding VM disk")
            end

            vm.cmd_nics.each do |nic|
                response=hyperv_command(host, nic)
                error_if(response, "Error adding VM NIC")
            end

            context=vm.cmd_context
            response=hyperv_command(host, context) if context
            error_if(response, "Error adding VM context")

            response=hyperv_command(host, vm.cmd_start_vm)
            error_if(response, "Error starting VM")

            vm.name
        end

        def poll(host, identifier)
            data=hyperv_command(host, "Get-VMState #{identifier}")
            error_if(data, "Error getting VM information")

            text=data[:stdout]

            # 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)
            response=hyperv_command(host, "Stop-VM #{identifier} -Force -Wait")
            error_if(response, "Error canceling VM")
            sleep 1
            delete(host, identifier)
        end

        def shutdown(host, identifier)
            response=cancel(host, identifier)
            error_if(response, "Error canceling VM")
        end

        def delete(host, identifier)
            response=hyperv_command(host, "Remove-VM #{identifier} -Force")
            error_if(response, "Error deleting VM")

            del_disks(host, identifier)
        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)
            response=ps(host,
                "Get-WmiObject win32_processor -ComputerName #{host} | "<<
                "select LoadPercentage,NumberOfCores | fl")
            error_if(response, "Error getting host info")

            text=response[:stdout]||""

            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=ps(host,
                "(Get-WmiObject -Class Win32_ComputerSystem "<<
                "-ComputerName #{host}).TotalPhysicalMemory")
            error_if(total, "Error getting host info")

            free=ps(host,
                "(Get-WmiObject -Class Win32_OperatingSystem "<<
                "-ComputerName #{host}).FreePhysicalMemory")
            error_if(free, "Error getting host info")

            {
                :total_mem => (total[:stdout]||nil).strip.to_i/1024,
                :free_mem => (free[:stdout]||nil).strip.to_i
            }
        end

        def del_disks(host, name)
            num=name.split('-')[-1]
            template="<TEMPLATE><VMID>#{num}</VMID></TEMPLATE>"
            vm=Template.new(template, @vmdir, @local_vmdir)

            del_cmd=vm.cmd_del_disks

            STDERR.puts del_cmd

            if del_cmd
                result=ps(host, del_cmd)
                STDERR.puts "Error deleting VM images" if result[:exitcode]!=0
            end
        end

        def hyperv_command(host, command)
            # ps("#{command} -server #{host}")
            if @proxy
                server_flag=" -server #{host}"
            else
                server_flag=""
            end

            ps(host, "#{command}#{server_flag}")
        end

        def ps(host, command)
            result=winrm(host).powershell("#{command}\n")

            #STDERR.puts command
            #STDERR.puts result.inspect

            data={
                :stdout => '',
                :stderr => ''
            }

            result[:data].each do |item|
                [:stdout, :stderr].each do |key|
                    if item[key]
                        data[key]<<item[key]
                        break
                    end
                end
            end

            data[:exitcode]=result[:exitcode]

            data
        end

        def error_if(code, message)
            if code && code.kind_of?(Hash) && code[:exitcode] &&
                    code[:exitcode]!=0
                OpenNebula.error_message(message)
                STDERR.puts code.inspect
                exit(-1)
            end
        end

    end
end

conf_file=ONE_ETC_PATH+"/hyperv.conf"

if !File.exist?(conf_file)
    OpenNebula.error_message('Can not read hyperv config file')
    exit(-1)
end

conf=YAML.load(File.read(conf_file))

vm_dir=conf[:vmdir]
local_vmdir=conf[:local_vmdir]
proxy=conf[:proxy]

USER=conf[:user]
PASSWORD=conf[:password]

if !vm_dir
    STDERR.puts "vmdir not set in the configuration file"
    exit(-1)
end

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

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


