#!/usr/bin/perl
# vim:ts=4 sw=4 ft=perl et:

use strict;
use warnings;
use Getopt::Long;
use LWP::UserAgent;
use LWP::Simple;
use MogileFS::Client;

# Written by Justin Brehm (jbrehm@icontact.com) - 05/07/2008 14:09:00
# 
# check_mogile.pl - Nagios Plugin to check the functionality of MogileFS
# Requirements - mysql-client 

my %ERRORS=('OK'=>0, 'WARNING'=>1, 'CRITICAL'=>2, 'UNKNOWN'=>3);
my %opts=('trackers'=> 'localhost:7001', 'domain'=>'default', 'dbhost'=> 'mysqldb001.example', 'db' => 'mogilefs', 'dbuser' => 'mogilefs', 'dbpass' => 'lolwut');

my $host = `uname -n`;
$host =~ s/[\r\n]//g;
my $test_key = "test_" . $host . "_" . time();

my $test_file = "/etc/mogilefs/mogilefs.conf"; # file that we are going to inject/extract 
unless ( -e $test_file ) {
    print "MOGILEFS UNKNOWN - File $test_file doesn't exist.  Please create and try again\n"; 
    exit($ERRORS{"UNKNOWN"});
}

my $mysql = '/usr/bin/mysql'; # specify path to mysql client binary
unless ( -X $mysql ) {
    print "MOGILEFS UNKNOWN - $mysql isn't executible or doesn't exist\n";
    exit($ERRORS{"UNKNOWN"});
}

abortWithUsage() unless
    GetOptions(
        't|trackers=s'    => \$opts{trackers},
        'd|domain=s'      => \$opts{domain},
        'host|dbhost=s'      => \$opts{dbhost},
        'db=s'          => \$opts{db},
        'u|dbuser=s'      => \$opts{dbuser},
        'p|dbpass=s'      => \$opts{dbpass},
        );

abortWithUsage() unless 
    ($opts{trackers} && $opts{domain} && $opts{db} && $opts{dbuser} && $opts{dbpass} && $opts{dbhost});

# Verify that the MySQL database is available first
my $mysql_check = `$mysql --user=$opts{dbuser} --password=$opts{dbpass} --host=$opts{dbhost} $opts{db} -e 'select 1' 2>/dev/null`;
if (!$mysql_check) {
    print "MOGILEFS CRITICAL - MogileFS database $opts{db} is unavailable\n";
    exit($ERRORS{"CRITICAL"});
}

my $mogfs = get_mogfs();
my $in_bytes = inject();
my $out_bytes = extract();
cleanup(); # Delete the file from MogileFS

if ($in_bytes eq $out_bytes) {
    print "MOGILEFS OK - Key $test_key inject/extract success\n";
    exit($ERRORS{"OK"});
} else {
    error("Store size: $in_bytes\nRetrieved size: $out_bytes");
}

sub abortWithUsage {
    print "check_mogile.pl -- Checks the functionality of MogileFS for iContact\n\n";
    print "Usage:\t\$ check_mogile.pl --trackers=127.0.0.1:7001 --domain=publisher [--dbhost=ral-sql002.colo] [--db=mogilefs] [--dbuser=icontact] -[-dbpass=secret]\n\n";
    print "\ttrackers\t: MogileFS trackers and port (127.0.0.1:7001)\n";
    print "\tdomain\t\t: MogileFS domain\n";
    print "\tdbhost\t\t: Hostname/IP of MySQL database\n";
    print "\tdb\t\t: The MogileFS database\n";
    print "\tdbuser\t\t: User with SELECT privileges on database\n";
    print "\tdbpass\t\t: password for user\n";
    exit($ERRORS{"UNKNOWN"});
}

sub error {
    my $err = shift() || "ERROR: no error message provided!";
    my $mogerr = undef;
    if ($mogerr = $mogfs->errstr) {
        $mogerr =~ s/^\s+//;
        $mogerr =~ s/\s+$//;
    }
    my $syserr = undef;
    
    if ($@) {
        $syserr = $@;
        $syserr =~ s/[\r\n]+$//;
    }
    
    print "CRITICAL - ";
    print "$err\n";
    print "MogileFS backend error message: $mogerr\n" if $mogerr;
    print "System error message: $@\n" if $syserr;
    exit($ERRORS{"CRITICAL"});
}

sub get_mogfs {
    my @trackerinput = split(/\s*,\s*/, $opts{trackers});
    my @trackers;
    my %pref_ip;

    foreach my $tracker (@trackerinput) {
        if ($tracker =~ m!(.+)/(.+):(\d+)!) {
            $pref_ip{$2} = $1;
            push @trackers, "$2:$3";
        } else {
            push @trackers, $tracker;
        }
    }
    
    my $handle = MogileFS::Client->new(
            domain => $opts{domain},
            hosts => \@trackers,
            timeout => 3,
            )
        or error("Could not initialize MogileFS");
    $handle->set_pref_ip(\%pref_ip);
    return $handle;
}

# Injects into the default class.  If we want to inject to a specific class
# we'll have to modify this.
sub inject {
    my $src = $test_file;
    my $key = $test_key;

    my $bytes = $mogfs->store_file($key, undef, $src) 
        or error("Injection test failed");

    return $bytes;
}

# Only returns that file size that is being extracted.  To get the full
# file you'd need to return $extractfile instead
sub extract {
    my $key = $test_key;
    my $file;
    
    $file->{type} = 'file';
    $file->{maxnum} = 1;
    $file->{parts}->{1} = {
        paths => [ grep { defined $_ } $mogfs->get_paths($key) ],
    };
    
    unless (scalar(@{$file->{parts}->{1}->{paths}})) {
        error("Extraction test failed");
    }
    
    my $extractfile;
    foreach my $path (@{$file->{parts}->{1}->{paths} || []}) {
        $extractfile .= get($path);
    }
    
    my $size = length $extractfile;
    return $size
}

sub cleanup {
    my $key = $test_key;

    unless ( $mogfs->delete($key) ) {
        error("Cannot delete $key");
    }
}