#!/usr/bin/ruby -w
# -*- ruby -*-

#
# dsck - Simple disk space checker
#
# Copyright: Tatsuki Sugiura <sugi@nemui.org>
# License: GPL-3
#

=begin

= NAME

((:dsck:)) - simple disk space checker

= SYNOPSIS

((%dsck [-Vhv] [-c config | global_threshold | /mount/point threshold]...%))

= DESCRIPTION

== Options

If both "-s" and "-r" not specified, set "-se" automatically.

: -V, --version
  Show version
: -h, --help
  Show this message
: -v, --verbose
  Verbose output. Multiple -v to increase verbosity
: -q, --quiet
  Quiet output. Multiple -q to decrease verbosity
: -s, --summary
  Show summary report
: -r, --report
  Show detail report
: -e, --errors-only
  Suppress "-r" and "-s" when nothing to alert
: -c ((|conf|)), --config=((|conf|))
  Specify config file (arguments prevail over config)

== Threshold Format

If argument is matched regex (({/[0-9]+([BKMGTb%])?/})) or
"nocheck", it's considerd as threshold.

Suffix character was consider as ((*B*))=byte, ((*K*))=KiloByte, ((*M*))=Megabyte,
((*G*))=GigaByte, ((*T*))=TeraByte, ((*P*))=PetaByte and ((*%*))=Percent.

When string has no suffix character, longer than ((*3*)) will
be considered ((*KiloByte*)), if not ((*Percent*)).

This script alert when use% is ((*over*)) threshold of percentage or
filesystem available size is ((*under*)) threshold of Byte and blocks.

== ConfigFile Format

Same as command line arguments.

Especially, after "#" until end-of-line is considerd as comment.

= Example

((%dsck -serv "95%" / 50M /tmp 5 /var/spool/squid nocheck%))

This means
* Verbose output (+1)
* Show summary and detail when error happened
* Global threshold is 95%, "/" is 50MByte, "/tmp" is 5%
* "/var/spool/squid" will not be checked

=end


require 'getoptlong'
require 'ostruct'

$MYNAME  = File.basename($0)
$VERSION = "3.0.0"

##
## classes
##
class DiskSpaceCheker < Hash
  class NonExistMountPoint < StandardError; end
  class InvalidThreshold   < StandardError; end
  DF_COLS = %w(device size used avail use_per mount)

  def initialize (global_threshold = "90%")
    @mounts = read_mounts
    @global_threshold = thres_normalizer(global_threshold)
  end

  def read_mounts
    mounts_dev = {}
    lno = 0
    `df -k`.each_line do |line|
      lno += 1
      lno == 1 and next
      cols = line.split(/\s+/)
      dev = cols.first
      dev == "" and
        cols[0] = dev = mounts_dev.keys.last
      mounts_dev[dev] ||= OpenStruct.new
      DF_COLS.each_with_index do |c, i|
        v = case c
            when "size", "avail", "used"
              cols[i].to_i * 1024
            else
              cols[i]
            end
        mounts_dev[dev].send("#{c}=", v)
      end
    end
    mounts = {}
    mounts_dev.each {|dev, m|
      mounts[m.mount] = m
    }
    mounts
  end

  def []=(key, val)
    unless @mounts.has_key?(key)
      raise NonExistMountPoint
    end
    val = thres_normalizer(val)
    super
  end

  def [](mount)
    self.has_key?(mount) ? super : @global_threshold
  end

  def need_alert?
    @mounts.each { |mount, minfo|
      if over?(mount)
	return true
      end
    }
    false
  end
  alias error? need_alert?

  def global_threshold=(thres)
    @global_threshold = thres_normalizer(thres)
  end

  def thres_normalizer(thres)
    if thres !~ /^\d+[BKMGTP%]?$/ && thres != "nocheck"
      raise InvalidThreshold
    elsif thres !~ /\D/
      if thres.length < 3
	thres = thres + "%"
      else
	thres = thres + "K"
      end
    end
    return thres
  end

  def over?(mount)
    unless @mounts.has_key?(mount)
      raise NonExistMountPoint, mount
    end
    minfo = @mounts[mount]
    thres = self[mount]
    if thres == "nocheck"
      return false
    elsif thres =~ /^(\d+)%$/
      thres = (minfo.size * (100 - $1.to_f) / 100).to_i
    elsif thres =~ /^(\d+)([BKMGTP])$/
      thres = $1.to_i * (1 << (10 * %w(B K M G T P).index($2).to_i))
    end
    minfo.avail <= thres
  end

  def summary
    ret = Array.new
    @mounts.each { |m, info|
      if over?(m)
	dev = info.device
	ret.push("Warning: #{dev} (#{m}) " +
		 (self[m] =~ /%$/ ? "is used over" : "less than") +
		 " #{self[m]}")
      end
    }
    ret.size > 0 ? ret : ["All good."]
  end

  def report
    ret = [%w(Device Mounted Size Avail Use% Thres Status)]
    @mounts.each { |mount, minfo|
      @mounts.has_key?(mount) or next
      use_per = ((minfo.size - minfo.avail) * 100 / minfo.size).to_i # NOTE: Do NOT use minfo.used due to reserved block
      ret.push [minfo.device, mount,
                human_size(minfo.size), human_size(minfo.avail),
                "#{use_per}%", self[mount],
                over?(mount) ? "=!ALERT!=" : "OK"]
    }
    ret.map{ |r|
      sprintf(%Q[#{ r[0].length < 15   ? "%-15s " : "%s\n#{' '*16}" }], r.shift) +
	sprintf(%Q[#{ r[0].length < 10 ? "%-10s " : "%s\n#{' '*27}" }], r.shift) +
	sprintf("%10s %10s %5s %8s %10s", *r)
    }
  end

  def human_size(size)
    suffixes = %w(K M G T P)
    suffixes.reverse.each_with_index do |s, i|
      unit = 1 << (10 * (suffixes.length - i))
      size < unit and next
      return sprintf("%0.1f#{s}", size / unit.to_f)
    end
    return size.to_s
  end

  attr_reader :global_threshold, :mounts
end


if File.basename(__FILE__) == File.basename($0)

  $echo_level = 3

  ##
  ## functions
  ##

  def version
    puts "#{$MYNAME} version #{$VERSION}.   Copyright 2002, Tatsuki Sugiura <sugi@nemui.org>"
  end

  def usage(io = $stdout)
    io.print <<EOU
Usage: #{$MYNAME} [-Vhvqsre] [-c config|global_threshold|/mount/point threshold]...
Options: (default="-se")
  -V,      --version      show version
  -h,      --help         show this message
  -v,      --verbose      verbose output. multiple -v for increase verbosity
  -q,      --quiet        quiet output. multiple -q for decrease verbosity
  -s,      --summary      show summary
  -r,      --report       show detail report
  -e,      --errors-only  suppress "-r" and "-s" when nothing to alert
  -c conf, --config=conf  specify config file (arguments prevail over config)
EOU
  end

  def msg(level = 3, *msg)
    # level:
    #  1  fatal error message
    #  2  warning
    #  3  notice, normal message
    #  4  debug output
    #  5- more...
    if $echo_level < level
      return true
    end
    (level < 3 ? $stderr : $stdout).print msg.join
  end

  ##
  ## main
  ##

  opt = Hash.new

  optparser = GetoptLong.new
  optparser.set_options(
			["--verbose", "-v",     GetoptLong::NO_ARGUMENT],
			["--quiet", "-q",       GetoptLong::NO_ARGUMENT],
			["--help", "-h",        GetoptLong::NO_ARGUMENT],
			["--version", "-V",     GetoptLong::NO_ARGUMENT],
			["--summary", "-s",     GetoptLong::NO_ARGUMENT],
			["--report", "-r",      GetoptLong::NO_ARGUMENT],
			["--errors-only", "-e", GetoptLong::NO_ARGUMENT],
			["--config", "-c",      GetoptLong::REQUIRED_ARGUMENT]
			)

  begin
    optparser.each_option do |name, arg|
      if name == "--verbose"
	$echo_level += 1
      elsif name == "--quiet"
	$echo_level -= 1
      else
	opt[name.sub(/^--/, '')] = arg
      end
    end
  rescue
    usage($stderr)
    exit(1)
  end

  if opt["version"]
    version
    exit(0)
  end

  if opt["help"]
    usage
    exit(0)
  end

  dsck = DiskSpaceCheker.new
  args = Array.new

  # read config
  if opt.has_key?("config")
    begin
      File.open(opt["config"]) { |fh|
	fh.each { |line|
	  line.sub!(/#.*/,"")
	  line.scan(/(\S+)/) { |part|
	    args.push(*part)
	  }
	}
      }
    rescue
      msg 2, "Error: Can't open config (#{$!}). Ignored.\n"
    end
  end

  args = args + ARGV
  while arg = args.shift
    if arg =~ /^\d+[BKMGTP%]?$/ || arg == "nocheck"
      dsck.global_threshold = arg
    else
      mount = arg
      threshold = args.shift
      begin
	dsck[mount] = threshold
      rescue DiskSpaceCheker::InvalidThreshold
	msg 2, "Error: Invalid threshold '#{threshold}' for #{mount}. Ignored.\n"
      rescue DiskSpaceCheker::NonExistMountPoint
	msg 2, "Warning: Mount point #{mount} is not found. Ignored.\n"
      end
    end
  end

  msg 5, "checking thresholds...\n"
  dsck.mounts.each {|m, minfo|
    msg 5, " #{m}\t#{dsck[m]}\n"
  }

  # set default option
  if !opt.has_key?("summary") && !opt.has_key?("report")
    opt["summary"] = opt["errors-only"] = true
  end

  if !opt["errors-only"] || dsck.need_alert?
    if opt["summary"]
      msg 4, "=== Problem summary ===\n"
      puts dsck.summary
    end

    if opt["report"]
      msg 4, "=== Current Status ===\n"
      puts dsck.report
    end
  end
  dsck.error? ? exit(1) : exit(0)
end
