#!/usr/bin/env ruby

# returns the mem usage of a given process
def plist(psname)
  counter = 0
  %x{ps h -o rss,size,vsize,pcpu -u puppet}.each do |ps|
    rss,size,vsize,cpu = ps.split
    counter += 1
    puts "puppetmasterd_#{counter}.value #{rss}"

  end
  return
end

# reports how many clients compiled in the last 5 minutes
# repotrs how many unique clients compiled since the begining of the day
# report the average compilation time for all clients in the last 5 minutes.
def phaselog(logfile)
  count,avg,day_count_unique,day_count = 0 ,0 ,0, 0
  t = Time.now
        today = t.strftime("^%b ") + " ?" + t.day.to_s
  hour = today + t.strftime(" %H:")
  m = t.min.to_i
  last5m = ""
  6.times do |i|
    last5m += hour
    last5m += "0" if (m-i) < 10
    last5m += (m-i).to_s
    last5m += "|" unless i==5
  end
  hosts = Array.new
  regexp = ".* for (.*) in (.*) seconds"

  File.open(logfile).grep(/#{today}/).grep(/Compiled configuration|Compiled catalog/).each do |line|
    case line
      when /#{last5m}/ then
        if line =~ /#{regexp}/
          avg += $2.to_f
          count += 1
          unless hosts.include?($1)
            hosts << $1
          end
        end
      when /#{regexp}/ then
        day_count += 1
        unless hosts.include?($1)
          hosts << $1
          day_count_unique += 1
        end
    end
  end
  puts "avg_compile.value #{count == 0 ? 0 : (avg / count).to_s[0..3]}"
  puts "last5m_count.value #{count}"
  puts "last24h_unique_count.value #{day_count_unique}"
end

sys_file = '/etc/default/puppetmaster'
sys_file = '/etc/sysconfig/puppetmaster' if !File.exists?(sys_file) && File.exists?('/etc/sysconfig/puppetmaster')
logfile = ENV['puppet_logfile'] || '/var/log/daemon.log'
facts_dir = ENV['facts_dir'] || '/var/lib/puppet/yaml/facts/'

case ARGV[0]
  when 'config'
    case $0
    when /puppetmaster_memory/
      puts "graph_title puppetmaster memory usage"
      puts "graph_vlabel memory"
      # find out how many mongrel process we have - if any
      File.open(sys_files) do |line|
        @pm_process = line.grep(/(PUPPETMASTERS|PUPPETMASTER_PORTS)/).to_s.split('=')[1].to_i
      end
      if @pm_process > 0
        @pm_process.times do |i|
          puts "puppetmasterd_#{i+1}.label puppetmasterd #{i+1}"
        end
      else
        puts "puppetmaster.label puppetmasterd"
      end
    when /puppet_clients/
      puts "graph_title puppet clients usage"
      puts "graph_vlabel clients"
      puts "known_clients.label Known Clients"
      puts "avg_compile.label Average configuration compile"
      puts "last5m_count.label Clients in the last 5 minutes"
      puts "last24h_unique_count.label unique clients in the last 24 hours"
    end
    puts "graph_category puppet"
    exit 0
  when 'autoconf'
    case $0
      when /puppet_mem/,/puppet_clients/
        if  File.exists?(sys_file) && File.exists?(logfile) && File.directory?(facts_dir)
          puts "yes"
        else
          puts "no"
        end
      else
        puts "no"
      exit 0
    end
  else
    plist("'ruby /usr/sbin/puppetmasterd'") if $0 =~ /puppet_mem$/
    if $0 =~ /puppet_clients$/
      puts "known_clients.value #{Dir.entries(facts_dir).size-2}"
      phaselog(logfile)
    end
end