require 'set'
require 'socket'

class NagiosPlugin < Plugin
    Config.register Config::StringValue.new('nagios.listen_addr',
      :default => "localhost", :requires_restart => true,
      :desc => "Address to listen on for Nagios notifications.")
    Config.register Config::IntegerValue.new('nagios.listen_port',
      :default => "5667", :requires_restart => true,
      :validate => Proc.new {|v| v > 0},
      :desc => "Port to listen on for Nagios notifications.")
    Config.register Config::ArrayValue.new('nagios.alert_channels',
      :default => [], 
      :desc => "List of channels or users to send notifications to.")
  
  def initialize
    super
    @notifications = "on"
    @ignore_alert_hosts = Set.new []
    listen_addr = (@bot.config['nagios.listen_addr']).to_s
    listen_port = (@bot.config['nagios.listen_port']).to_i
    nagios_alert_channels = @bot.config['nagios.alert_channels']
    @nsca_thread = nil
    if(!@nsca_thread.nil?)
      @nsca_thread.raise ThreadRestartError, "Restarting thread"
    end
    @nsca_thread = Thread.new do
      begin
        Thread::abort_on_exception=true
        nsca_socket = TCPserver.open(listen_addr, listen_port)
        _nsca_thread_startup nsca_socket, nagios_alert_channels
        debug "Normal shutdown of nsca thread"
        nsca_socket.shutdown
        nsca_socket.close
      rescue ThreadRestartError
        debug "Closing nsca thread"
        nsca_socket.shutdown
        nsca_socket.close
      end
    end # Thread.new
  end

  class ThreadRestartError < ThreadError
  end

  def _nsca_thread_startup(nsca_socket, nagios_alert_channels)
    debug nsca_socket.addr.inspect
    #while true
    # Thread.start(nsca_socket.accept) do |s|
    #   debug "#{s} is accepted"
    #   while s.gets
    #     s.write($_)
    #     debug $_
    #   end
    #   print "#{s} is gone"
    #   s.close
    # end
    #end
    while true
      debug "Looping #{nsca_socket.addr.inspect}"
      Thread.start(nsca_socket.accept) do |sk|
        Thread::abort_on_exception=true
        debug "Recieving #{sk.inspect}"
        while sk.gets
          s = $_
          sk.write "Got message: #{s}"
          debug  "Got message: #{s}\n"
          if s.length > 0 and @notifications == 'on'
            alert_data = s.split("#|#")
            alert_data.collect! {|x| x.strip! }
            debug "Data=#{alert_data.inspect}"
            msg = _nsca_format_message(alert_data)
            debug "Output: #{msg}"
            if(not @ignore_alert_hosts.member?(alert_data[2]))
              nagios_alert_channels.each { |a| 
                debug "Sending to: #{a}"
                @bot.notice(a,"(nagios) #{msg}")
              }
            else
                debug "Ignoring message for host: #{alert_data[2]}"
            end
          end
        end
        debug "Closing #{sk.inspect}"
        sk.close
      end # Thread.start
      debug "End of loop\n"
    end # while true
    debug "Completed nsca thread\n"
  end

##-jeeves/#gentoo-infra- (nagios) **RECOVERY** Service: procs-total | Host: warbler | State: OK | Info: PROCS OK: 125 processes | Date: Tue Oct 31 06:29:20 UTC 2006
## command_line /bin/echo "SERVICE-ALERT #|# $NOTIFICATIONTYPE$ #|# $HOSTNAME$ #|# $SERVICEDESC$ #|# $SERVICESTATE$ #|# $OUTPUT$ #|# $DATETIME$" | /usr/nagios/libexec/send_nsca 10.0.0.1 -p 5668 -c /etc/nagios/send_nsca.cfg
## command_line /bin/echo "HOST-ALERT #|# $NOTIFICATIONTYPE$ #|# $HOSTNAME$ #|# $HOSTSTATE$ #|# $OUTPUT$ #|# $DATETIME$" | /usr/nagios/libexec/send_nsca 10.0.0.1 -p 5668 -c /etc/nagios/send_nsca.cfg

  def _nsca_format_message(msgarr,sep='|')
    alerttype = msgarr[0]
    notificationtype = msgarr[1]
    hostname = msgarr[2]
    if(alerttype == 'SERVICE-ALERT')
      servicename = msgarr[3]
      state = msgarr[4]
      output = msgarr[5]
      datetime = msgarr[6]
      msgextra = "Service: #{servicename} #{sep} "
    elsif(alerttype == 'HOST-ALERT')
      state = msgarr[3]
      output = msgarr[4]
      datetime = msgarr[5]
      msgextra = ""
    end
    msg = "**#{notificationtype}** "
    msg << msgextra
    msg << "Host: #{hostname} #{sep} "
    msg << "State: #{state} #{sep} "
    msg << "Info: #{output} #{sep} "
    msg << "Date: #{datetime}"
    return msg
  end
    

  def toggle_global(m,params)
    debug params.inspect
    debug m
    @notifications = params[:mode]
    m.reply "Turning #{@notifications} notifications" 
  end

  def status(m,params)
    debug params.inspect
    debug m
    m.reply "Notifications are #{@notifications}, and the following hosts are off: #{@ignore_alert_hosts.to_a.sort.join(' ')}"
  end

  def toggle_host(m,params)
    debug params.inspect
    mode = params[:mode]
    host = params[:host]
    if(mode == 'on')
      @ignore_alert_hosts.delete(host)
    elsif(mode == 'off')
      @ignore_alert_hosts << host
    end
    m.reply "Turning #{mode} notifications for host #{host}"
  end

  def cleanup
    debug "Cleaning up thread and socket\n"
    @nsca_thread.exit
  end

  def help(cmd, topic="")
    case topic
      when "off"
        return "nagios off: turns off nagios listener"
      when "on"
        return "nagios on: turns on nagios listener"
      when "status"
        return "nagios status: prints nagios listener status"
      when /^host/
        return "nagios host <on|off>: controls nagios listener for a single host"
      else
        return "#{name}: on, off, status, host <on|off>"
    end
  end
end

plugin = NagiosPlugin.new
plugin.map 'nagios :mode', 
  :action => 'toggle_global',
  :requirements => {:mode => /^on$|^off$/}
plugin.map 'nagios status', 
  :action => 'status'
plugin.map 'nagios :host :mode', 
  :action => 'toggle_host',
  :requirements => {:mode => /^(on|off)$/, :host => /^\w+$/}

# end
# vim: ts=2 sw=2 et: