#! /usr/local/bin/ruby

#  Copyright (c) 2006 Frédéric Senault.  All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. Neither the name of Frédéric Senault or any contributors may be
#    used to endorse or promote products derived from this software
#    without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.

$LOAD_PATH << File.dirname($0)

require 'syslog'
require 'common'

# The part that extracts the contents of the performance data and stores it
# in rrd databases.
class TMain < Main
	# Finds all the collecters plugins, and store them in the global config
	# (tester only).
	def loadcollecters()
		Dir.glob($conf[:collpath]) do |f|
			begin
				require "#{f}"
			rescue Exception => e
				puts e
			end
		end
	end

	# Finds all the tester plugins, instanciate then, sort them by order, and
	# store the objects in the global config (tester only).
	def loadtesters()
		Dir.glob($conf[:testpath]) do |f|
			begin
				require "#{f}"
			rescue Exception => e
				puts e
			end
		end
		$conf[:tests] = []
		Tester.globtests.sort! { |a, b| a.order <=> b.order }
		Tester.globtests.each { |t| $conf[:tests].push(t.new) }
	end

	# The main loop ; parses data sources, and present data to the appropriate
	# plugins (tester only).
	def parsefile()
		allrrds = {}
		allperf = []
		Collecter.globcolls.each do |oc|
			c = oc.new
			c.each do |nd|
				ok = false
				$conf[:tests].each do |tst|
					begin
						tst.register(nd) do |id|
							name = tst.rrdname(nd, id)
							rrd = allrrds[name]
							if(rrd.nil?) then
								rrd = RRDStruct.new(name)
								rrd.attachtest(tst, id)
								allrrds[name] = rrd
							end
							rrd.attachdata(nd)
							dolog(:info, tst.class.to_s + ' took ' + \
								nd.short)
							ok = true
						end
						break if(ok)
					rescue Exception => e
						dolog(:error, e)
					end
				end
				dolog(:info, 'Nobody took ' + \
					nd.short) unless(ok)
			end
		end
		begin
			allrrds.each_value do |rrd|
				rrd.initrrdstruct()
				rrd.matchdata()
				dolog(:notice, 'Inserting data : ' + rrd.name + ', ' + \
					rrd.counttimes.to_s + 't, ' + \
					rrd.countds.to_s + 's, ' + \
					rrd.countdp.to_s + 'd')
				rrd.writerrd()
				rrd.clear()
			end
		rescue Exception => e
			dolog(:error, e)
		end
	end
end

# The main collecter class ; must be inherited by the child plugins to work.
class Collecter
	include Enumerable
	@@gcolls = []

	# This is triggered when the class is inherited ; it allows the class to maintain
	# a class variable to store it's children.
	def Collecter.inherited(c)
		@@gcolls.push(c)
	end

	# Gives the list of plugins classes.
	def Collecter.globcolls
		@@gcolls
	end

	# The plugin must be Enumerable.
	def each
		raise "Child must implement each."
	end
end

# This class implements the basic recommended data for the collecter plugins.
# Of course, it can be extended by the child plugins and the testers if needed.
class NGraphData
	# Unix timestamp for the performance data
	attr_reader :time

	# Hostname.
	attr_reader :host

	# The performance data itself ; its format is free - as long as the
	# tester plugins can make something out of it.
	attr_reader :perfdata

	# Transforms the timestamp into a Time class.
	def ptime
		 ptime = Time.at(@time)
	end

	# Sort id of the data, used for the logs.
	def short()
		@host
	end
end

# The main tester class ; it has to be inherited by all the plugins.  The flow is :
# 1. The order class property is queried to order the plugins.
# 2. The object is instanciated.
# 3. The data is offered to the plugins through the register method.
# 4. If the register method yielded one or more identifiers, the rrdname method
#    is called with the data and the id.
# 5. The initrrdstruct is then called with a new RRDStruct object to populate and the id.
# 6. The getdata is called with the RRDStruct, the data, and the id.
class Tester
	@@gtests = []

	# This is triggered when the class is inherited ; it allows the class to maintain
	# a class variable to store it's children.
	def Tester.inherited(c)
		@@gtests.push(c) unless(c.to_s == 'ServiceTest')
	end

	# The list of tester plugins.
	def Tester.globtests
		@@gtests
	end

	# The order in which the plugin will be interrogated (lowest order first).
	def self.order
		999
	end

	# This is called with the data to examine.  Must be implemented, and must
	# yield one or more identifiers if it wants the.  The ndata object is
	# totally opaque for the tester system.
	def register(ndata)
		raise "register must be implemented"
	end

	# This is called to generate a filename for the rrd files.  By default,
	# the name is generated based on the id of the source and the hostname.
	def rrdname(ndata, id)
		ndata.host.normalize + '_' + id.to_s.normalize
	end

	# Called to populate the RRD file and datastructure.  The rrd parameter
	# is a RRDStruct object.  This function must be implemented.
	def initrrdstruct(rrd, id)
		raise "inittrdstruct must be implemented"
	end

	# The last called function, to store the data itself.  The rdd parameter
	# is the RRDStruct object initialized before, and it also takes the identifier.
	def getdata(rrd, ndata, id)
		raise "getdata must be implemented"
	end
end

# Compatibilty shim.
class ServiceTest < Tester
end

# The main datastructure of the tester system.  It represents one rrd file, but
# usually multiple values at diverse measure times.
# The method useable from the plugins are :
# 1. addds : from the initrrdstruct method in the plugin.
# 2. addtime : from the getdata method in the plugin.
# 3. data[]= : from the getdata method, after addtime.
class RRDStruct
	include Enumerable

	# A hash table of the attached RRDData objects.
	attr_accessor :data

	# The basic interval for the highest grained precision - usually defined
	# at a system level with the $conf[:step] value.
	attr_accessor :step
	
	# The name if the rrd ; will be expanded to $conf[:rrdspath] + '/' + name + '.rrd'.
	attr_accessor :name

	# The odest stored data.
	attr_reader :inittime

	# The list of RRA objects linked to this file.  Usually defined at a system
	# level with the $conf[:rras] value.  Can be overridden via the addrra function.
	attr_reader :rras

	# The RRDStruct is created empty with only a name.  It will be filled by the
	# different methods of the plugins.
	def initialize(name)
		@name = name
		@fname = $conf[:rrdspath] + '/' + name + '.rrd'
		@tst = {}
		@rras = nil
		@data = {}
		@values = {}
		@nd = nil
	end

	# Adds a data source to the file ; returns a RRDData object.  It must be called
	# in the Tester#initrrdstruct method.  See the RRDData#initialize method for
	# the parameters.
	def addds(name, type = nil, heartbeat = nil, min = nil, max = nil)
		@data[name] = RRDData.new(name, type, heartbeat, min, max)
	end

	# Called internally to link a line of data to the RRDStruct.
	def attachdata(nd)
		@nd = [] if(@nd.nil?)
		@nd.push(nd)
	end

	# Must be called when a new timestamp is received.  Returns a hash table with
	# the keys of the datasources as a key, and 'U' (undefined) for a balue.
	def addtime(time)
		if(@values[time].nil?) then
			@values[time] = {}
			@data.each do |k, v|
				@values[time][k] = 'U'
			end
		end
		@values[time]
	end
	
	# Adds a customized rra ; see RRA#new for the parameters ; if this method
	# is called once, the default $conf[:rras] is not used anymore.
	def addrra(funct = :AVERAGE, xff = 0.5, steps = nil, \
		   rows = nil, title = nil, start = nil)
		@rras ||= []
		@rras.push(RRA.new(funct, xff, steps, rows, title, start))
	end

	# Clears all data from this measure.
	def clear
		@values = {}
	end

	# Attaches a plugin to the RRDStruct when it takes the data.
	def attachtest(tst, id)
		@tst[id] = tst
	end

	# Called internally to instanciate the rrdstructs.
	def initrrdstruct
		@tst.each do |id, t|
			t.initrrdstruct(self, id)
		end
	end

	# Simple counter of the different data measure timestamps.  Used for logging.
	def counttimes
		@values.keys.length
	end

	# Counter of the datasources.  Used for logging.
	def countds
		@data.keys.length
	end

	# Counter of all the non unknown points.  Used for logging.
	def countdp
		m = 0
		@values.each_key do |k|
			@values[k].each_value do |v|
				m += 1 unless(v.nil? || v == 'U')
			end
		end
		m
	end

	# The data points can be enumerated.
	def each
		@data.each_key.sort do |t|
			yield(t, @data[t])
		end
	end
	
	# Called internally to collect the data.
	def matchdata
		@nd.each do |s|
			@tst.each do |id, t|
				t.getdata(self, s, id)
			end
		end
	end

	# Called internally to create the rrd files themselves.
	def createfile
		unless(File.exists?(@fname) || @values.keys.length == 0) then
			@inittime = @values.keys.sort[0] - 1
			@step = $conf[:step]
			@rras = $conf[:rras] if(@rras.nil?)
			system_d($conf[:rrdtool] + ' ' + to_rra)
			if($?.exitstatus != 0) then
				puts "Création #{@name} : #{$?.exitstatus}"
			end
		end
	end
	
	# Called internally to update the datapoints.
	def writerrd
		createfile
		if(@data.keys.length > 0 && @values.keys.length > 0) then
			c = to_rru
			system_d($conf[:rrdtool] + ' ' + c)
			if($?.exitstatus != 0) then
				puts "Update #{@name} : #{$?.exitstatus}"
			end
		end
	end

	# Called internally to generate the <i>rrdtool create</i> command line.
	def to_rra()
		c = []
		c.push("create #{@fname}")
		c.push("-b #{@inittime}")
		c.push("-s #{@step}")
		@data.each_value do |v|
			c.push(v.to_rra)
		end
		@rras.each do |v|
			c.push(v.to_rra)
		end
		c.join(' ')
	end

	# Called internally to generate the <i>rrdtool update</i> command line.
	def to_rru()
		c = []
		c.push("update #{@fname}")
		tm = []
		@data.keys.sort.each do |k|
			tm.push(k)
		end
		c.push("-t #{tm.join(':')}")
		@values.keys.sort.each do |t|
			tv = t.to_s
			@values[t].keys.sort.each do |k|
				tv += ':' + @values[t][k].to_s
			end
			c.push(tv)
		end
		c.join(' ')
	end
end

# Object returned from the addds method, defines the different data to store in the rrd.
class RRDData

	# The name of the datasource.
	attr_reader :name
	
	# The type of the datasource, one of <tt>:GAUGE</tt>, <tt>:COUNTER</tt>,
	# <tt>:ABSOLUTE</tt> or <tt>:DERIVE</tt>.
	attr_reader :type

	# A new datasource.  The name is the key of the datasource in the file, the type
	# is one of the fundamental rrdtools types (by default, <tt>:GAUGE</tt>).
	# The heartbeat is the minimum interval at which the data must be updated.  The
	# min and max values define the validity of the inserted data.  If unspecified,
	# any values will be allowed.
	def initialize(name, type = nil, heartbeat = nil, min = nil, max = nil)
		@name = name
		@type = ( type.nil? ? :GAUGE : type )
		@value = 'U'
		@heartbeat = ( heartbeat.nil? ? $conf[:heartbeat] : heartbeat )
		@min = ( min.nil? ? 'U' : min )
		@max = ( max.nil? ? 'U' : max )
	end

	# Used internally to generate the command line for <tt>rrdtool create</tt>.
	def to_rra
		[ 'DS', @name, @type.to_s, @heartbeat.to_s, @min.to_s, @max.to_s ].join(':')
	end
end

main = TMain.new()
main.readparam()
main.loadcollecters()
main.loadtesters()
main.parsefile()