#!/usr/bin/ruby

require 'yaml'


#
#  This class either a) returns the output of running commands, or
# b) returns the output from a faked command
#
class CommandWrapper

  def CommandWrapper.run_command( str )

    if ( ENV['TEST'] && ENV['TEST_PREFIX'] )
      #
      #  Count the number of commands we've executed so far.
      #
      count = ENV['TEST_COUNT'] || '0'
      count = count.to_i
      ENV['TEST_COUNT'] = (count + 1).to_s

      #
      #  Read the output from the faked file.
      #
      #    $prefix/$count.cmd
      #
      file = "#{ENV['TEST_PREFIX']}/#{count}.cmd"
      File.open( file, "r" ).readlines().join()
    else
      `#{str}`
    end
  end
end

#
# Class to handle mapping of kernel block device names to disc-free (df)
# output.
#
class BMMonProbeDiscFree
  BIN_DF = '/bin/df'
  PROC_DISKSTATS = '/proc/diskstats'

  attr_writer :bin_df, :proc_diskstats

  def initialize
    @bin_df = nil
    @proc_diskstats = nil
  end

  #
  # This generates a lookup table of [major][minor] to kernel name.
  #
  def lookup
    lookup =  Hash.new { |h, k| h[k] = {} }

    proc_diskstats.each do |line|
      line = line.chomp.split(/\s+/)
      lookup[line[1].to_i][line[2].to_i] = line[3]
    end

    lookup
  end

  #
  # This gives the disc usage and mount points each filesystem described by
  # bin_df
  #
  def disc_free(ignored_fs = [], _show_dummy = false)
    do_parse_df(bin_df('-i'), ignored_fs, do_parse_df(bin_df('-k'), ignored_fs))
  end

  protected

  #
  # This returns the contents of /proc/diskstats or the value of
  # @proc_diskstats.  This latter function is useful for testing.
  #
  def proc_diskstats
    return @proc_diskstats unless @proc_diskstats.nil?
    fail Errno::ENOENT, PROC_DISKSTATS unless File.exist?(PROC_DISKSTATS)
    File.open(PROC_DISKSTATS) { |fh| fh.readlines }
  end

  #
  # This returns output from df (called with the arguments -a -k --portability
  # --local --print-type) or the value of @bin_df.  This latter function is
  # useful for testing.
  #
  def bin_df(unit = '-k')
    return @bin_df unless @bin_df.nil?
    fail Errno::ENOENT, BIN_DF unless File.exist?(BIN_DF)
    fail Errno::ENOEXEC, "#{BIN_DF} is not executable. This is not normal!" unless File.executable?(BIN_DF)

    CommandWrapper.run_command("#{BIN_DF} -a #{unit} --portability --local --print-type 2>&1")
  end

  def do_parse_df(df, ignored_fs = [], disc_free =  Hash.new { |h, k| h[k] = Hash.new(0) })
    units = nil

    df.split("\n").each do |line|
      data = line.chomp.split(/\s+/)

      if data[1] == 'Type'
        if 'Inodes' == data[2]
          units = 'inodes'
        else
          units = 'kbytes'
        end
        next
      end

      #
      # Work out the mountpoint
      #
      mountpoint = File.expand_path(data[6])
      mountpoint = 'root' if mountpoint == '/'

      # Ignore proc/sys/dev mountpoints
      next if mountpoint =~ /^\/(proc|sys|dev)(\b|\/)/

      #
      # Miss any ignored fs types.
      #
      next if ignored_fs.include?(data[1])

      disc_free[mountpoint]['device'] = data[0]
      disc_free[mountpoint]['fs_type'] = data[1]
      disc_free[mountpoint]["total_#{units}"] = data[2].to_i
      disc_free[mountpoint]["used_#{units}"] = data[3].to_i
      disc_free[mountpoint]["available_#{units}"] = data[4].to_i
      disc_free[mountpoint]["use_percent_#{units}"] = data[5].to_i
      disc_free[mountpoint]['mount_point'] = data[6]
    end

    disc_free
  end
end

#
#  Entry point to our code.
#
if __FILE__ == $PROGRAM_NAME

  def verbose(str)
    STDERR.puts(str)
  end

  #
  #  File types we ignore
  #
  ignored_fs  = %w(iso9660 fuse.gvfs-fuse-daemon debugfs aufs)

  #
  #  The results of checking the disks
  #
  disk_result = BMMonProbeDiscFree.new.disc_free(ignored_fs)

  #
  #  The alerts we'll raise, if any.
  #
  to_raise = []

  #
  # Check for ignored filesystems
  #
  ignored = []
  if File.exist?('/etc/disk-free.ignore')
    ignored = File.readlines('/etc/disk-free.ignore').collect(&:chomp)
  end

  #
  # Reset the minimum threshold if there is a file.
  #
  min = 95
  if File.exist?('/etc/disk-free.threshold')
    new_min = File.readlines('/etc/disk-free.threshold').first
    unless new_min.nil?
      min = new_min.chomp.to_i
    end
  end

  verbose "Alerting on partitions with usage greater than #{min}%"

  #
  # This hash is used to save writing the same code out twice, once for disc
  # usage, and once for inode usage.
  #
  [%w(disc use_percent_kbytes),
   %w(inode use_percent_inodes)].each do |what, field|
    disk_result.each do |mount, details|
      unless details.key?(field)
        verbose "Ignoring #{mount} as it has no usage information"
        next
      end

      if ignored.any? { |m| mount =~ /^#{m}/ }
        verbose "Ignoring #{mount} as it matches an entry in /etc/disk-free.ignore"
        next
      end

      if %w(proc none sysfs).include?(details['fs_type'])
        verbose "Ignoring #{mount} as its filesystem is one of proc, sysfs, or none."
        next
      end

      h = {}
      h[:id] =  "#{what.downcase}-#{mount}"
      h[:summary] =  "Mount point #{mount} has #{100 - details[field]}% #{what} free - #{details[field]}% #{what} used"
      h[:detail] = " * #{mount}\n" + details.sort.collect { |k, v| " * #{k}: #{v}" }.join("\n")

      if details[field] > min
        verbose "#{what} usage on #{mount} is #{details[field]}% -- sending alert"

        to_raise.push(h)
      else
        verbose "#{what} usage on #{mount} is #{details[field]}% -- not sending alert"
      end
    end
  end

  #
  #  Show any alerts we have.
  #
  puts YAML.dump(to_raise)
  exit 0
end
