#!/usr/bin/perl -w
#############################################################################
# License:
#
# This software (hereafter referred to as "program") is free software;
#   you can redistribute it and/or modify it under the terms of the GNU General
#   Public License as published by the Free Software Foundation; either version
#   2 of the License, or (at your option) any later version.
# Note that when redistributing modified versions of this source code, you
#   must ensure that this disclaimer and the above coder's names are included
#   VERBATIM in the modified code.
#
# Disclaimer:
# This program is provided with no warranty of any kind, either expressed or
#   implied.  It is the responsibility of the user (you) to fully research and
#   comprehend the usage of this program.  As with any tool, it can be misused,
#   either intentionally (you're a vandal) or unintentionally (you're a moron).
#   THE AUTHOR(S) IS(ARE) NOT RESPONSIBLE FOR ANYTHING YOU DO WITH THIS PROGRAM
#   or anything that happens because of your use (or misuse) of this program,
#   including but not limited to anything you, your lawyers, or anyone else
#   can dream up.  And now, a relevant quote directly from the GPL:
#
#                           NO WARRANTY
#
#  11. BECAUSE THE PROGRAM IS LICENSED FREE OF CHARGE, THERE IS NO WARRANTY
# FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW.  EXCEPT WHEN
# OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES
# PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED
# OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.  THE ENTIRE RISK AS
# TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU.  SHOULD THE
# PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING,
# REPAIR OR CORRECTION.
#
#   ---
#
# Whee, that was fun, wasn't it?  Now let's all get together and think happy
#   thoughts - and remember, the same class of people who made all the above
#   legal spaghetti necessary are the same ones who are stripping your rights
#   away with DVD CCA, DMCA, stupid software patents, and retarded legal
#   challenges to everything we've come to hold dear about our internet.
#   So enjoy your dwindling "freedom" while you can, because 1984 is coming
#   sooner than you think.  :[
#
#############################################################################
# 
# Purpose:
#     A script designed to be run at incremental periods (via crond) which 
#     will monitor disk throughput, and syslog that information.
# 
# 
# Technical Description:
#     Reads data out of /proc/stat and generates output.  It uses a temporary
#     file called /tmp/monitorDiskThroughput.XXX
# 
# 
# Changelog:
#     01/20/2004 - v1.2
#         - Require a -r option to run.
#         - Added a -l option for logging to a file.
#         - Textual changes to most error messages.
#         - Try opening /proc/stat up to 10 times.
#         - Don't use any system commands to retreive data.
#         
#     09/30/2002 - v1.0-rc5
#         - Lots of textual changes
#         - Added the printmsg() function
# 
# 
####################################################################################################################################
use strict;

## Global Variable(s)
my %conf = (
    "programName"          => $0,
    "version"              => '1.20',
    "authorName"           => 'Brandon Zehm',                    ## Information about the author or owner of this script.
    "authorEmail"          => 'caspian@dotconf.net',
    "hostname"             => $ENV{'HOSTNAME'},                  ## Used in printmsg() for all output.
    
    "syslog"               => 1,                                 ## Syslog messages by default (--stdout disables this)
    "syslogProgram"        => '/usr/local/bin/syslog.pl',        ## This could also be the standard "logger" utility.
    
    "tempFileThroughput"   => '/tmp/monitorDiskThroughput',      ## Temp file used to store counter values in between runs.
    
    "run"                  => 0,                                 ## Whether we should run script or not.
    "debug"                => 0,                                 ## Change this with the -v command line option.
    "stdout"               => 0,
    "facility"             => 'USER',
    "priority"             => 'INFO',
    "logging"              => '',                                ## If this is true the printmsg function prints to the log file
    "logFile"              => '',                                ## If this is specified (form the command line via -l) this file will be used for logging.
    
);
$conf{"programName"} =~ s/(.)*[\/,\\]//;                         ## Remove path from programName







#############################
##                          ##
##      MAIN PROGRAM         ## 
##                          ##
#############################

## Initialize
initialize();

## Process Command Line
processCommandLine();

## Open the log file if we need to
if ($conf{'logFile'}) {
    if ($conf{'logging'}) {
        $conf{'logging'} = 0;
        close LOGFILE;
    }
    if (openLogFile($conf{'logFile'})) { quit("OS-ERROR => Opening the log file [$conf{'logFile'}] returned the error: $!", 1); }
}


## Make sure it's not a 2.6 kernel
if (-f "/proc/diskstats") {
    quit("OS-ERROR => The Linux 2.6 kernel is not currently supported by this script.", 1);
}

## Monitor disk throughput for all drives
my @devices = ();
for (my $counter = 0; $counter <= 10; $counter++, sleep(10)) {
    quit("OS-ERROR => 10 consecutive errors while trying to read /proc/stat.  Last error: $!", 1) if ($counter >= 10);
    open(STAT, "< /proc/stat") or next;
    my @lines = <STAT>;
    @lines = grep(/disk_io/,@lines);
    if (scalar(@lines) < 1) {
        quit("OS-NOTICE => /proc/stat does not contain disk_io information.  No statistics can be gathered.", 1);
    }
    @devices = split(/ /, $lines[0]);
    close STAT;
    chomp @devices;
    last if ($devices[0]);
}
shift(@devices);

foreach my $device (@devices) {
    if ( ($device) and ($device =~ /\d/) ) {
        chomp $device;
        printmsg("OS-DEBUG => Device string is: $device", 1);
        $device =~ s/\(//; $device =~ s/\).*//;
        printmsg("OS-DEBUG => Device string is now: $device",1);
        $conf{'message'} = ("OS-$conf{'priority'} => " . monitorThroughput($device, $conf{'tempFileThroughput'}) );
        
        ## Print/syslog this message
        printmsg($conf{'message'},0);
    }
}


## Quit
quit("",0);
































######################################################################
## Function:    help ()
##
## Description: For all those newbies ;) 
##              Prints a help message and exits the program.
## 
######################################################################
sub help {
print <<EOM;

$conf{'programName'}-$conf{'version'} by $conf{'authorName'} <$conf{'authorEmail'}>

Summary:
  Monitors disk throughput. The system syslog is used for all output
  unless the --stdout option is used.

Usage:  $conf{'programName'} [options]
  
  Required:
    -r                        run script

  Optional:
    --stdout                  print messages to STDOUT rather than the syslog
    --facility=[0-11]         syslog facility (1/USER is used by default)
    -l <logfile>              enable logging to the specified file
    -v                        verbosity - use multiple times for greater effect

EOM
exit(1);
}










######################################################################
##  Function: initialize ()
##  
##  Does all the script startup jibberish.
##  
######################################################################
sub initialize {

    ## Set STDOUT to flush immediatly after each print  
    $| = 1;
    
    ## Intercept signals
    $SIG{'QUIT'}  = sub { quit("WARNING => EXITING: Received SIG$_[0]", 1); };
    $SIG{'INT'}   = sub { quit("WARNING => EXITING: Received SIG$_[0]", 1); };
    $SIG{'KILL'}  = sub { quit("WARNING => EXITING: Received SIG$_[0]", 1); };
    $SIG{'TERM'}  = sub { quit("WARNING => EXITING: Received SIG$_[0]", 1); };
    $SIG{'HUP'}   = sub { quit("WARNING => EXITING: Received SIG$_[0]", 1); };
    $SIG{'ALRM'}  = sub { quit("WARNING => EXITING: Received SIG$_[0]", 1); };
    
    ## Fixup $conf{'hostname'}
    if ($conf{'hostname'}) {
        $conf{'hostname'} = $conf{'hostname'};
        $conf{'hostname'} =~ s/\..*$//;
    }
    else {
        $conf{'hostname'} = "unknown";
    }
  
    ## Return 0 errors
    return(0);
}









######################################################################
##  Function: processCommandLine ()
##  
##  Processes command line storing important data in global var %conf
##  
######################################################################
sub processCommandLine {
  
    ############################
    ##  Process command line  ##
    ############################
  
    my @ARGS = @ARGV;
    my $numargv = @ARGS;
    my $counter = 0;
    for ($counter = 0; $counter < $numargv; $counter++) {

        if ($ARGS[$counter] =~ s/^--stdout//i) {             ## stdout ##
            $conf{'stdout'} = 1;
            $conf{'syslog'} = 0;
        }
        
        elsif ($ARGS[$counter] =~ s/^--facility=//i) {       ## Facility ##
            $conf{'facility'} = $';
        }
        
        elsif ($ARGS[$counter] =~ /^-r$/) {                  ## Run
            $conf{'run'} = 1;
        }
            
        elsif ($ARGS[$counter] =~ /^-l$/) {                  ## Log File ##
            $counter++;
            $conf{'logFile'} = $ARGS[$counter];
        }
        
        elsif ($ARGS[$counter] =~ s/^-v+//i) {               ## Verbosity ##
            $conf{'debug'} += (length($&) - 1);
        }
        
        elsif ($ARGS[$counter] =~ /^-h$|help/i) {            ## Help ##
            help();
        }
        
        else {                                               ## Unknown Option ##
            printmsg("ERROR => The option [$ARGS[$counter]] is not understood! Try --help.", 0);
        }
        
    }
  
    ## If the user didn't use a -r print the help
    if ($conf{'run'} == 0) {
        help();
    }
  
    return(0);
}

















######################################################################
##  Function:    monitorThroughput (string $device, string $tempFile)
##  
##  Description: Device will be a string like: "8,1" (meaning scsi disk 1)
##               You can get a list of current devices by looking at /proc/stat
##               
##
##  Example:     
######################################################################
sub monitorThroughput {
    my %incoming = ();
    (
      $incoming{'device'},
      $incoming{'tempFile'}
    ) = @_;
    
    unless ($incoming{'device'}) {
        return("monitorThroughput() - No device specified!");
    }
    
    $incoming{'tempFile'} .= ".$incoming{'device'}";
    my $readSectors = "";
    my $writeSectors = "";
    my $lastReadSectors;
    my $lastWriteSectors;

    my $sectorSize = 512;
    my $currentTime = time();
    my $lastTime;
    my $transferRate;
    my $null;
    my $stats;
    
    ## Get raw data from /proc/stats
    for (my $counter = 0; $counter <= 10; $counter++, sleep(10)) {
        quit("OS-ERROR => 10 consecutive errors while trying to read /proc/stat.  Last error: $!", 1) if ($counter >= 10);
        open(STAT, "/proc/stat") or next;
        my @lines = <STAT>;
        close STAT;
        chomp @lines;
        @lines = grep(/^disk_io:/,@lines);
        if (scalar(@lines) < 1) {
            quit("OS-NOTICE => /proc/stat does not contain disk_io information.  No statistics can be gathered.", 1);
        }
        ## Grab the info for the device we're interested in
        ($stats) = grep(/\($incoming{'device'}\)/, split(/ /, $lines[0]));
        
        ## Remove junk we don't need
        $stats =~ s/\($incoming{'device'}\):\(//;
        $stats =~ s/\)$//;
        
        ## Get current IN and OUT bytes
        ($null,$null,$readSectors,$null,$writeSectors) = split(/,/, $stats);
        
        ## Exit the loop if we got our data
        last if (($readSectors =~ /\d/) and ($writeSectors =~ /\d/));
    }
    
    
    
    ## Get data from the temp file, and write the new values
    if ( -f $incoming{'tempFile'} ) {
        
        open(FILE, $incoming{'tempFile'}) or quit("OS-ERROR => Error while opening the file [$incoming{'tempFile'}].  The error was: $!", 1);
        
        ($lastTime, $lastReadSectors, $lastWriteSectors) = split(/:/, <FILE>);
        
        ## Write new values to tempFile
        open (FILE, ">$incoming{'tempFile'}") or quit("OS-ERROR => Error while opening the file [$incoming{'tempFile'}].  The error was: $!", 1);
        print FILE "$currentTime:$readSectors:$writeSectors\n";
        close FILE;
        
        ## Deal with the 4GB counter limitation of the Linux kernel.
        if ($lastReadSectors > $readSectors) { $readSectors += (4294967296 - $lastReadSectors); $lastReadSectors = 0; };
        if ($lastWriteSectors > $writeSectors) { $writeSectors += (4294967296 - $lastWriteSectors); $lastWriteSectors = 0; };
        
        ## Here we re-assign these three variables to be the difference from last time... I know it's evil.
        $currentTime = ($currentTime - $lastTime);
        $readSectors = ($readSectors - $lastReadSectors);
        $writeSectors = ($writeSectors - $lastWriteSectors);
        
        ## Calculate the difference between the last read and now, and calculate that into KBytes/second
        $transferRate = (((($readSectors + $writeSectors) * $sectorSize) / $currentTime) / 1024);
        
        
    }
    
    ## Write the current data to the temp file if its a new file
    else {
        printmsg("NOTICE => The file [$incoming{'tempFile'}] does not exist.  Creating it now with mode [0600].", 0); 
        open (FILE, ">$incoming{'tempFile'}") or quit("OS-ERROR => Error while creating the file [$incoming{'tempFile'}].  The error was: $!", 1);
        
        print FILE time() . ":$readSectors:$writeSectors\n";
        close FILE;
        chmod (0600, $incoming{'tempFile'});
        
        $currentTime = 0;
        $readSectors = 0;
        $writeSectors = 0;
        $transferRate = 0;
    }
    
    
    
    
    
    return(sprintf("Period: $currentTime seconds  Device: $incoming{'device'}  Sectors Read: $readSectors  Sectors Written: $writeSectors  Transfer Rate: %.02f KB/s", $transferRate));
    
}
















###############################################################################################
##  Function:    printmsg (string $message, int $level)
##
##  Description: Handles all messages - 
##               Depending on the state of the program it will log
##               messages to a log file, print them to STDOUT or both.
##               
##
##  Input:       $message          A message to be printed, logged, etc.
##               $level            The debug level of the message. If not defined 0
##                                 will be assumed.  0 is considered a normal message, 
##                                 1 and higher is considered a debug message.
##  
##  Output:      Prints to STDOUT, to LOGFILE, both, or none depending 
##               on the state of the program and the debug level specified.
##  
##  Example:     printmsg("ERROR => The file could not be opened!", 0);
###############################################################################################
sub printmsg {
    ## Assign incoming parameters to variables
    my ( $message, $level ) = @_;
    
    ## Make sure input is sane
    $level = 0 if (!defined($level));
    
    ## Continue only if the debug level of the program is >= message debug level.
    if ($conf{'debug'} >= $level) {
        
        ## Get the date in the format: Dec 03 11:14:04
        my ($sec, $min, $hour, $mday, $mon) = localtime();
        $mon = ('Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec')[$mon];
        my $date = sprintf("%s %02d %02d:%02d:%02d", $mon, $mday, $hour, $min, $sec);
    
        ## Syslog the message is needed
        if ($conf{'syslog'}) {
            system( qq{$conf{'syslogProgram'} --facility=$conf{'facility'} --priority=$conf{'priority'} "$conf{'programName'}: $message" });
        }
        
        ## Print to STDOUT always if debugging is enabled, or if conf{stdout} is true.
        if ( ($conf{'debug'} >= 1) or ($conf{'stdout'} == 1) ) {
            print "$date $conf{'hostname'} $conf{'programName'}: $message\n";
        }
        
        ## Print to the log file if $conf{'logging'} is true
        if ($conf{'logging'}) {
            print LOGFILE "$date $conf{'hostname'} $conf{'programName'}: $message\n";
        }
        
    }
    
    ## Return 0 errors
    return(0);
}














###############################################################################################
## FUNCTION:    
##   openLogFile ( $filename )
## 
## 
## DESCRIPTION: 
##   Opens the file $filename and attaches it to the filehandle "LOGFILE".  Returns 0 on success
##   and non-zero on failure.  Error codes are listed below, and the error message gets set in
##   global variable $!.
##   
##   
## Example: 
##   openFile ("/var/log/scanAlert.log");
##
###############################################################################################
sub openLogFile {
    ## Get the incoming filename
    my $filename = $_[0];
    
    ## Make sure our file exists, and if the file doesn't exist then create it
    if ( ! -f $filename ) {
        printmsg("NOTICE => The file [$filename] does not exist.  Creating it now with mode [0600].", 0);
        open (LOGFILE, ">>$filename");
        close LOGFILE;
        chmod (0600, $filename);
    }
    
    ## Now open the file and attach it to a filehandle
    open (LOGFILE,">>$filename") or return (1);
    
    ## Put the file into non-buffering mode
    select LOGFILE;
    $| = 1;
    select STDOUT;
    
    ## Tell the rest of the program that we can log now
    $conf{'logging'} = "yes";
    
    ## Return success
    return(0);
}











######################################################################
##  Function:    quit (string $message, int $errorLevel)
##  
##  Description: Exits the program, optionally printing $message.  It 
##               returns an exit error level of $errorLevel to the 
##               system  (0 means no errors, and is assumed if empty.)
##
##  Example:     quit("Exiting program normally", 0);
######################################################################
sub quit {
    my %incoming = ();
    (
        $incoming{'message'},
        $incoming{'errorLevel'}
    ) = @_;
    $incoming{'errorLevel'} = 0 if (!defined($incoming{'errorLevel'}));
    
    
    ## Print exit message
    if ($incoming{'message'}) { 
        printmsg($incoming{'message'}, 0);
    }
    
    ## Exit
    exit($incoming{'errorLevel'});
}






