# vim: ts=2 sts=2 sw=2 et si
# $Id: checkservice.rb 2153 2012-12-29 13:10:58Z aremizov $

require 'yaml'

module MH
  module CheckService
    attr_reader :cfg, :service_state

    class PerfData
      attr_reader :value

      def initialize(data)
        if data.kind_of? Hash
          @label = data[:label]
          @value = %w{value warning critical min max}.map { |k| data[k.to_sym] }.join ';'
        else
          @label = 'data'
          @value = data
        end
        @as_string = "'#{@label.to_s.gsub(/'=/, '')}'=#{@value}"
      end

      def to_s
        @as_string
      end
    end

    class ServiceStatus < Exception
      EXIT_CODE = {
        :ok       => 0,
        :warning  => 1,
        :critical => 2,
        :unknown  => 3
      }

      WEIGHT = {
        :ok       => 0,
        :unknown  => 1,
        :warning  => 2,
        :critical => 3
      }

      attr_reader :status, :perf_data, :exit_code

      def initialize(status, message, perf_data = nil)
        if EXIT_CODE.include? status
          @status, @message = status, message
        elsif EXIT_CODE.index(status)
          @status, @message = EXIT_CODE.index(status), message
        else
          @status, @message = :unknown, "Unknown status #{status} (original message is \"#{message}\")"
        end
        @exit_code = EXIT_CODE[@status]
        @perf_data = if perf_data and not perf_data.empty?
          if perf_data.kind_of? Array
            perf_data.flatten
          else
            PerfData.new perf_data
          end
        end
      end

      def <=>(service_status)
        WEIGHT[service_status.status] <=> WEIGHT[@status]
      end

      def message
        if @perf_data.respond_to? :value
          @message % [@perf_data.value]
        else
          @message
        end
      end

      def to_s
        final_message = if @perf_data
          perf_data = if @perf_data.kind_of? Array
            @perf_data.map { |d| d.to_s }.join ' '
          else
            @perf_data.to_s
          end
          "#{message}|#{perf_data}"
        else
          message
        end

        "#{status.to_s.upcase} #{final_message}"
      end
    end

    class State
      def initialize(filename)
        @filename = filename
        @current_state = Hash.new
      end

      def load
        @common_config = {
          'history_length' => 10
        }
        if not File.zero? @filename and File.readable? @filename
          @all_states = YAML.load_file @filename
          unless @all_states.kind_of? Array
            raise ServiceStatus.new :unknown, "Can't parse state file #{@filename}"
          end
        else
          @all_states = []
        end
      end

      def <<(current_state)
        @all_states.unshift({ :time => Time.new, :state => current_state})
        save 
      end

      def save
        while @all_states.size > @common_config['history_length'] do
          @all_states.pop
        end
        File.open(@filename, 'w', 0700) { |f| YAML.dump(@all_states, f) }
      end

      def get(n = 1)
        @all_states[n]
      end

      def set(k,v)
        self << {k => v}
      end
    
      def current_state
        @all_states.first[:state]
      end

      def difference(axis, item = nil)
        return if @all_states.length < 2
        if item
          @all_states[0][:state][item][axis] - @all_states[1][:state][item][axis]
        else
          @all_states[0][:state][axis] - @all_states[1][:state][axis]
        end
      end
    
      def derivative(axis, item = nil)
        return if @all_states.length < 2
        timedelta = @all_states[0][:time] - @all_states[1][:time]
        return if timedelta == 0
        difference(axis, item) / timedelta
      end
      
    end 

    def check_service(service_name)
      read_config service_name
      @service_state = State.new File.join(@common_config['statedir'], "#{service_name}.state.yaml")

      lock_file = File.join @common_config['lockdir'], "#{service_name}.lock"
      service_status = if lock_script(lock_file)
        begin
          @service_state.load

          states = [yield].flatten.compact.sort
          status :unknown, 'no data collected' if states.empty?

          general_status = states.first.status
          general_message = states .
            select { |state| state.status == general_status } .
            map { |state| state.message } .
            join ', '
          general_perf_data = states.map { |state| state.perf_data }.compact

          status general_status, general_message, general_perf_data
        rescue ServiceStatus => service_status
          service_status
        ensure
          unlock_script(lock_file)
        end
      else
        ServiceStatus.new :unknown, 'script execution is locked'
      end
      puts service_status.to_s
      exit service_status.exit_code
    end

    def status_message(message, perf_data = nil)
      @status_message = message
      @perf_data = perf_data
    end

    def status(status, message = nil, perf_data = nil)
      status_message(message, perf_data) if message
      raise ServiceStatus.new status, @status_message, @perf_data
    end

    def status_from_external_command(command)
      message = %x{#{command}}.to_s.strip
      status $?.exitstatus, message
    end

    def status_from_file(filename, max_age)
      unless File.readable? filename
        status :unknown, "Can't read file #{filename}"
      end
      if Time.now - File.stat(filename).ctime > max_age
        status :unknown, "Status file #{filename} is older than #{max_age} seconds"
      end

      status_code, message = File.open(filename).read.split(' ', 2)
      unless status_code.kind_of? String
        status_code = 'unknown'
        message = "Can't recognize status from '#{filename}'"
      end
      status status_code.downcase.to_sym, message
    end

    def qx(commandline)
      result = %x{#{commandline}}
      status :unknown, "'#{commandline}' invocation error" unless $?.exitstatus == 0
      result
    end

    def defaults(hash)
      @defaults = hash
    end

    def read_config(service_name)
      def locate_config(filename)
        ['/usr/local/etc', '/etc'].each do |path|
          config_file = File.join(path, filename)
          return config_file if File.exists? config_file
        end
        nil
      end

      @common_config = {
        'lockdir'   => '/tmp',
        'statedir' => '/tmp',
      }
      @cfg = @defaults.clone

      config_file = locate_config('checkservice.conf')
      return unless config_file

      sections = File.open(config_file) { |f| YAML::load(f) }
      if sections.kind_of? Hash
        @common_config.update sections['common'] if sections.include? 'common'
        @cfg.update sections[service_name] if sections.include? service_name
      end
    end

    def lock_script(lock_file)
      if File.exists? lock_file
        %x{pgrep -F#{lock_file} -f #{$0}}
        if $?.exitstatus == 0
          puts "Lock file #{lock_file} found. Abort execution."
          throw :lock_found
        end
      end

      File.open(lock_file, 'w') { |f| f.write $$ }
    end

    def unlock_script(lock_file)
      begin
        File.unlink lock_file
      rescue Errno::ENOENT
      end
    end
  end
end

class Array
  def check_each
    self.map do |item|
      begin
        yield item
        status :unknown, "no status for \"#{item}\""
      rescue MH::CheckService::ServiceStatus => service_status
      end
      service_status
    end
  end
end