#!/usr/bin/env ruby

# -------------------------------------------------------------------------- #
# Copyright 2002-2017, OpenNebula Project, OpenNebula Systems                #
#                                                                            #
# 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.                                             #
#--------------------------------------------------------------------------- #

$: << File.join(File.dirname(__FILE__), '../lib')

require 'pp'
require 'rexml/document'
require 'base64'
require 'uri'

require "poll_common"

begin
    require 'rubygems'
    require 'json'

    JSON_LOADED = true
rescue LoadError
    JSON_LOADED = false
end

ENV['LANG']='C'
ENV['LC_ALL']='C'

################################################################################
#
#  KVM Monitor Module
#
################################################################################
module KVM
    # Constants for KVM operations
    CONF={
        :dominfo    => 'virsh --connect LIBVIRT_URI --readonly dominfo',
        :list       => 'virsh --connect LIBVIRT_URI --readonly list',
        :dumpxml    => 'virsh --connect LIBVIRT_URI --readonly dumpxml',
        :domifstat  => 'virsh --connect LIBVIRT_URI --readonly domifstat',
        :domblkstat => 'virsh --connect LIBVIRT_URI --readonly domblkstat',
        :top        => 'top -b -d2 -n 2 -p ',
        'LIBVIRT_URI' => 'qemu:///system'
    }

    # Execute a virsh command using the predefined command strings and URI
    # @param command [Symbol] as defined in the module CONF constant
    def self.virsh(command)
        CONF[command].gsub('LIBVIRT_URI', CONF['LIBVIRT_URI'])
    end

    # Get the information of a single VM. In case of error the VM is reported
    # as not found.
    # @param vm_id [String] with the VM information
    def self.get_vm_info(one_vm)
        dominfo = dom_info(one_vm)

        return { :state => '-' } if !dominfo

        psinfo = process_info(dominfo['UUID'])

        vm = Hash.new

        vm[:name] = one_vm
        vm[:pid]  = psinfo[1]

        cpu = get_cpu_info({one_vm => vm})

        resident_mem = psinfo[5].to_i
        max_mem      = dominfo['Max memory'].split(/\s+/).first.to_i

        values=Hash.new

        values[:state]  = get_state(dominfo['State'])
        values[:cpu]    = cpu[vm[:pid]] if cpu[vm[:pid]]
        values[:memory] = [resident_mem, max_mem].max

        xml = dump_xml(one_vm)

        values.merge!(get_interface_statistics(one_vm, xml))

        return values
    end

    # Gets the information of all VMs
    #
    # @return [Hash, nil] Hash with the VM information or nil in case of error
    def self.get_all_vm_info
        vms_info = Hash.new
        vms      = Hash.new

        text=`#{virsh(:list)}`

        return nil if $?.exitstatus != 0

        lines = text.split(/\n/)[2..-1]

        names = lines.map do |line|
            line.split(/\s+/).delete_if {|d| d.empty? }[1]
        end

        return vms_info if names.length == 0

        names.each do |vm|
            dominfo = dom_info(vm)

            if dominfo
                psinfo = process_info(dominfo['UUID'])

                info= Hash.new

                info[:dominfo] = dominfo
                info[:psinfo]  = psinfo
                info[:name]    = vm
                info[:pid]     = psinfo[1]

                vms[vm]=info
            end
        end

        cpu = get_cpu_info(vms)

        vms.each do |name, vm|
            ps_data = vm[:psinfo]
            dominfo = vm[:dominfo]

            resident_mem = ps_data[5].to_i
            max_mem      = dominfo['Max memory'].split(/\s+/).first.to_i

            values = Hash.new

            values[:state]  = get_state(dominfo['State'])
            values[:cpu]    = cpu[vm[:pid]] if cpu[vm[:pid]]
            values[:memory] = [resident_mem, max_mem].max

            xml = dump_xml(name)

            values.merge!(get_interface_statistics(name, xml))

            if !name.match(/^one-\d+/)
                uuid, template = xml_to_one(xml)
                values[:template] = Base64.encode64(template).delete("\n")
                values[:vm_name] = name
                vm[:name] = uuid
            end

            values.merge!(get_diskio_statistics(name, xml))

            vms_info[vm[:name]] = values
        end

        return vms_info
    end

    def self.number_of_processors
        %x{nproc}.to_i
    end

    def self.get_cpu_jiffies
        begin
            stat = File.read("/proc/stat")
        rescue
            return 0
        end

        jiffies = 0

        # skip cpu string and guest jiffies
        stat.lines.first.split(' ')[1..-3].each do |num|
            jiffies += num.to_i
        end

        jiffies
    end

    def self.get_process_jiffies(pid)
        begin
            stat = File.read("/proc/#{pid}/stat")
        rescue
            return 0
        end

        jiffies = 0

        data = stat.lines.first.split(' ')

        [13, 14, 15, 16].each do |col|
            jiffies += data[col].to_i
        end

        jiffies
    end

    # Gathers process information from a set of VMs.
    #   @param vms [Hash] of vms indexed by name. Value is a hash with :pid
    #   @return  [Hash] with ps information
    def self.get_cpu_info(vms)
        pids = vms.map {|name, vm| vm[:pid] }
        pids.compact!

        multiplier = number_of_processors * 100

        cpu = Hash.new

        start_cpu_jiffies = get_cpu_jiffies

        pids.each do |pid|
            cpu[pid] = get_process_jiffies(pid).to_f
        end

        sleep 1

        cpu_jiffies = get_cpu_jiffies - start_cpu_jiffies

        pids.each do |pid|
            cpu[pid] = ( get_process_jiffies(pid) - cpu[pid] ) / cpu_jiffies
            cpu[pid] = ( cpu[pid] * multiplier ).round(2)
        end

        cpu
    end

    # Process information for a KVM domain by its UUID
    #   @param uid [String] with user id
    #   @return [Array] of user processes
    def self.process_info(uuid)
        ps=`ps auxwww | grep -- '-uuid #{uuid}' | grep -v grep`
        ps.split(/\s+/)
    end

    # Gets the info of a domain by its id
    #   @param the ID of the VM as defined in libvirt
    #   @return [Hash] with the output of virsh dominfo, indexed by name (Id...)
    # Example execution of dominfo
    #   Id:             5
    #   Name:           one-6
    #   UUID:           06bc1876-fc6a-4dca-b41d-d7f2093b6b59
    #   OS Type:        hvm
    #   State:          running
    #   CPU(s):         1
    #   CPU time:       11.1s
    #   Max memory:     524288 KiB
    #   Used memory:    524288 KiB
    #   Persistent:     no
    #   Autostart:      disable
    #   Managed save:   no
    #   Security model: none
    #   Security DOI:   0
    def self.dom_info(vmid)
        text = `#{virsh(:dominfo)} #{vmid} 2>/dev/null`

        return nil if $?.exitstatus != 0

        lines = text.split(/\n/)
        hash  = Hash.new

        lines.map do |line|
            parts = line.split(/:\s+/)

            hash[parts[0]] = parts[1]
        end

        hash
    end

    # Get dumpxml output of a VM
    #   @param the ID of the VM as defined in libvirt
    #   @return [String] xml output of virsh dumpxml
    def self.dump_xml(vmid)
        `#{virsh(:dumpxml)} '#{vmid}' 2>/dev/null`
    end

    # Aggregate statics of all VM NICs
    #   @param the ID of the VM as defined in libvirt
    #   @param text [nil, String] dumpxml output or nil to execute dumpxml
    #   @return [Hash] with network stats, by name [symbol] :netrx, :nettx
    def self.get_interface_statistics(vmid, text = nil)
        text = dump_xml(vmid) if !text

        return {} if $?.exitstatus != 0

        doc = REXML::Document.new(text)

        interfaces = Array.new

        doc.elements.each('domain/devices/interface/target') do |ele|
            interfaces << ele.attributes["dev"]
        end

        return {} if interfaces.empty?

        values = Hash.new

        values[:netrx] = 0
        values[:nettx] = 0

        interfaces.each do |interface|
            text=`#{virsh(:domifstat)} #{vmid} #{interface} 2>/dev/null`

            next if $?.exitstatus != 0

            text.each_line do |line|
                columns = line.split(/\s+/)

                case columns[1]
                    when 'rx_bytes'
                        values[:netrx] += columns[2].to_i
                    when 'tx_bytes'
                        values[:nettx]+=columns[2].to_i
                    end
                end
        end

        values
    end

    def self.get_disk_names(vmid, text = nil)
        text = dump_xml(vmid) if !text

        doc=REXML::Document.new(text)
        disks = []
        doc.elements.each('domain/devices/disk/target') do |ele|
            disks << ele.attributes["dev"]
        end

        disks
    end

    def self.get_diskio_statistics(vmid, text = nil)
        disks=get_disk_names(vmid, text)

        if disks && !disks.empty?
            values={}
            values[:diskrdbytes]=0
            values[:diskwrbytes]=0
            values[:diskrdiops]=0
            values[:diskwriops]=0

            disks.each do |disk|
                text=`#{virsh(:domblkstat)} #{vmid} #{disk}`

                text.each_line do |line|
                    columns=line.split(/\s+/)
                    case columns[1]
                    when 'rd_bytes'
                        values[:diskrdbytes]+=columns[2].to_i
                    when 'wr_bytes'
                        values[:diskwrbytes]+=columns[2].to_i
                    when 'rd_req'
                        values[:diskrdiops]+=columns[2].to_i
                    when 'wr_req'
                        values[:diskwriops]+=columns[2].to_i
                    end
                end
            end

            values
        else
            {}
        end
    end

    # Translate libvirt state to Opennebula monitor state
    #  @param state [String] libvirt state
    #  @return [String] OpenNebula state
    #
    # Libvirt states for the guest are
    #  * 'running' state refers to guests which are currently active on a CPU.
    #  * 'blocked' not running or runnable (waiting on I/O or in a sleep mode).
    #  * 'paused' after virsh suspend.
    #  * 'shutdown' guest in the process of shutting down.
    #  * 'dying' the domain has not completely shutdown or crashed.
    #  * 'crashed' guests have failed while running and are no longer running.
    #
    def self.get_state(state)
        case state.gsub('-', '')
            when *%w{running blocked shutdown dying idle paused}
                'a'
            when 'crashed'
                'e'
            else
                '-'
        end
    end

    # Convert the output of dumpxml to an OpenNebula template
    #   @param xml [String] output of dumpxml
    #   @return [Array] uuid and OpenNebula template encoded in base64
    def self.xml_to_one(xml)
        doc = REXML::Document.new(xml)

        name = REXML::XPath.first(doc, '/domain/name').text
        uuid = REXML::XPath.first(doc, '/domain/uuid').text
        vcpu = REXML::XPath.first(doc, '/domain/vcpu').text
        memory = REXML::XPath.first(doc, '/domain/memory').text.to_i / 1024
        arch = REXML::XPath.first(doc, '/domain/os/type').attributes['arch']

=begin
        disks = []
        REXML::XPath.each(doc, '/domain/devices/disk') do |d|
            type = REXML::XPath.first(d, '//disk').attributes['type']
            driver = REXML::XPath.first(d, '//disk/driver').attributes['type']
            source = REXML::XPath.first(d, '//disk/source').attributes[type]
            target = REXML::XPath.first(d, '//disk/target').attributes['dev']

            disks << {
                :type => type,
                :driver => driver,
                :source => source,
                :target => target
            }
        end

        disks_txt = ''

        disks.each do |disk|
            disks_txt << "DISK=[\n"
            disks_txt << "  SOURCE=\"#{disk[:source]}\",\n"
            disks_txt << "  DRIVER=\"#{disk[:driver]}\",\n"
            disks_txt << "  TARGET=\"#{disk[:target]}\""
            disks_txt << "]\n"
        end


        interfaces = []
        REXML::XPath.each(doc,
                "/domain/devices/interface[@type='bridge']") do |i|
            mac = REXML::XPath.first(i, '//interface/mac').
                attributes['address']
            bridge = REXML::XPath.first(i, '//interface/source').
                attributes['bridge']
            model = REXML::XPath.first(i, '//interface/model').
                attributes['type']

            interfaces << {
                :mac => mac,
                :bridge => bridge,
                :model => model
            }
        end

        interfaces_txt = ''

        interfaces.each do |interface|
            interfaces_txt << "NIC=[\n"
            interfaces_txt << "  MAC=\"#{interface[:mac]}\",\n"
            interfaces_txt << "  BRIDGE=\"#{interface[:bridge]}\",\n"
            interfaces_txt << "  MODEL=\"#{interface[:model]}\""
            interfaces_txt << "]\n"
        end
=end

        spice = REXML::XPath.first(doc,
            "/domain/devices/graphics[@type='spice']")
        spice = spice.attributes['port'] if spice

        spice_txt = ''
        if spice
            spice_txt = %Q<GRAPHICS = [ TYPE="spice", PORT="#{spice}" ]>
        end

        vnc = REXML::XPath.first(doc, "/domain/devices/graphics[@type='vnc']")
        vnc = vnc.attributes['port'] if vnc

        vnc_txt = ''
        if vnc
            vnc_txt = %Q<GRAPHICS = [ TYPE="vnc", PORT="#{vnc}" ]>
        end


        feature_list = %w{acpi apic pae}
        features = []

        feature_list.each do |feature|
            if REXML::XPath.first(doc, "/domain/features/#{feature}")
                features << feature
            end
        end

        feat = []
        features.each do |feature|
            feat << %Q[  #{feature.upcase}="yes"]
        end

        features_txt = "FEATURES=[\n"
        features_txt << feat.join(",\n")
        features_txt << "]\n"


        template = <<EOT
NAME="#{name}"
CPU=#{vcpu}
VCPU=#{vcpu}
MEMORY=#{memory}
HYPERVISOR="kvm"
IMPORT_VM_ID="#{uuid}"
OS=[ARCH="#{arch}"]
#{features_txt}
#{spice_txt}
#{vnc_txt}
EOT

        return uuid, template
    end
end

################################################################################
# MAIN PROGRAM
################################################################################

hypervisor = KVM
file       = 'kvmrc'
vars       = %w{LIBVIRT_URI}

load_vars(hypervisor, file, vars)

vm_id = ARGV[0]

if vm_id == '-t'
    print_all_vm_template(hypervisor)
elsif vm_id
    print_one_vm_info(hypervisor, vm_id)
else
    print_all_vm_info(hypervisor)
end
