#!/usr/bin/perl
#
# Copyright (C) Nikhef 2011
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author:
#     Mischa Sall\'e <msalle@nikhef.nl>
#     NIKHEF Amsterdam, the Netherlands
#
########################################################################
#
# Nagios probe to test functioning of EES
#
# Nagios state can be one of the following:
#   OK:
#          EES is up and running and responds correctly to SAML2-XACML2-request
#   WARNING:
#          EES is running and responds, but either slowly or not with 200 OK
#   CRITICAL:
#          EES does not respond or not in time
#   UNKNOWN
#
########################################################################

# DEFAULTS
my $probeversion="0.2.2";

# Note the following defaults can be overridden using cmdline options
my $deftimeout=10;	# Overall timeout for probe
my $defwarning=1;	# When to warn about slow running
my $defhost="localhost";# Default hostname for EES service
my $defport=6217;	# Default portnumber for EES service

########################################################################
# Logging package
#   keeps internal log trace which can be dumped with dump_log 
########################################################################
package logger;
use strict;
use warnings;
{
    my $loglevel;
    my @logstring;

    # Constructor
    sub new	{
	my $classname=shift;
	my $self={}; bless $self;
	my $level=shift;
	if (defined $level) {
	    $self->set_loglevel($level);
	} else {
	    $loglevel=0;
	}
	return $self;
    }

    # Sets loglevel
    sub set_loglevel($) {
	my $self=shift;
	my $level=shift;
	$loglevel=$level;
    }

    # Logging function: log_func(priority, "logstring\n");
    sub log_func($@) {
	my $self=shift;
	my $prio=shift;
	return if ($prio > $loglevel);
	for my $line (@_)	{
	    push @logstring,$line;
	}
    }

    # Dumps log
    sub get_log(@)	{
	my $self=shift;
	foreach my $myentry ( @logstring )  {
	    print $myentry;
	}
    }
}

########################################################################
# Nagios status printing package
#   Can set and dump nagios status output
########################################################################
package nagstat;
{
    my $code;
    my $summary;
    my $perfdata;
    my @stat;

    # Constructor
    sub new()   {
	my $classname=shift;
	my $self={}; bless $self;
	$code=3; # Default status unknown
	$summary=undef;
	$perfdata=undef;
	@stat=("OK","WARNING","CRITICAL","UNKNOWN");
	return $self;
    }

    # Set nagios code (0-3) plus summary
    sub set_status($$)   {
	my $self=shift;
	if (!defined $summary)	{
	    $code=shift;
	    $summary=shift;
	}
    }

    # Set internal performance data
    sub set_perfdata($)   {
	my $self=shift;
	$perfdata=shift;
    }

    # Printout nagios status, summary and optionally performance data
    # return value is code (0-3)
    sub get_status	{
	if (!defined $summary)	{
	    $summary="unknown status";
	}
	if (defined $perfdata)    {
	    print $stat[$code].": ".$summary."|".$perfdata."\n";
	} else {
	    print $stat[$code].": ".$summary."\n";
	}
	return $code;
    }
}

########################################################################
# Inter process communication package for nagios probes
#   Starts alarm handler when receiving alarm which checks status of
#   probe, and terminates or kills it.
########################################################################
package probeipc;
{
    my $timeout;

    # Constructor: new(exitfunc,[kill time], [term time])
    sub new()   {
	my $classname=shift;
	my $self={}; bless $self;
	my $alarmhandler=shift
	    or die ($classname."::new() needs alarmhandler arg\n");
	my $inthandler=shift
	    or die ($classname."::new() needs inthandler arg\n");
	my $timeout=(shift or 10); # probe default timeout is 10
	$self->set_alarmhandler($alarmhandler);
	$self->set_inthandler($inthandler);
	$self->set_timeout($timeout);
	return $self;
    }

    # Sets time after which to send SIGKILL 
    sub set_timeout($)	{
	my $self=shift;
	$timeout=shift;
    }
    
    # Sets function to call when SIGALRM is caught
    sub set_alarmhandler($)	{
	my $self=shift;
	my $alarmhandler=shift;
	# \& for function reference, $ for stringvar
	$SIG{'ALRM'} = \&$alarmhandler;
    }

    # Sets function to call when SIGINT or SIGTERM is caught
    sub set_inthandler($)	{
	my $self=shift;
	my $inthandler=shift;
	$SIG{'INT'} = \&$inthandler;
	$SIG{'TERM'} = \&$inthandler;
    }

}

########################################################################
# Running main probe package
########################################################################
package main;
use strict;
use warnings;

use IO::Socket;
use Getopt::Long qw(:config no_ignore_case bundling);
use Time::HiRes qw(time alarm);

my $timeout;	# Total maximum runtime for probe
my $critical;	# Time after which to kill gLExec
my $warning;	# Time after which to warn about slow gLExec
my $host;	# EES hostname
my $port;	# EES portnumber
my $verbose;	# Verbosity level

my $sock;	# socket to EES

# Define different stages, such that e.g. the sighandlers know where we are
my %stages=(
    'presock'	=> 0,
    'sockopen'	=> 1,
    'datasent'	=> 2,
    'headrcvd'	=> 3,
    'resprcvd'	=> 4,
    'sockclosed'=> 5
);
my $stage=$stages{'presock'};	# Which state the socket is in

# Prints usage output
sub usage() {
    (my $name = $0) =~ s/.*\///;
    print <<EOHELP;
Usage: $name [options]

Options:
 -t|--timeout <timeout>          maximum runtime for probe, default: $deftimeout sec
 -c|--critical <timeout>         idem
 -w|--warning <time>             runtime after which to warn, default: $defwarning sec
 -H|--host <hostname>            hostname, default: $defhost
 -p|--port <portnumber>          port number, default: $defport
 -v|--verbose                    be more verbose, more -v means more verbosity
 -V|--version                    print version
 -h|--help                       show this helptext
EOHELP
    exit 0;
}

# Prints short usage output (oneline)
sub shortusage()	{
    (my $name = $0) =~ s/.*\///;
    print <<EOHELP;
Usage: $name [options]
EOHELP
}

# Prints probe version
sub version()	{
    (my $name = $0) =~ s/.*\///;
    print <<EOHELP;
$name version: $probeversion
EOHELP
}

# Parses command line options and sets global variables
sub getopts()	{
    my $version;
    my $help;
    my $shorthelp;

    $timeout=$deftimeout;
    $warning=$defwarning;
    $host=$defhost;
    $port=$defport;
    GetOptions(
	"t|timeout=f" => \$timeout,
	"c|critical=f" => \$timeout,
	"w|warning=f" => \$warning,
	"H|host=s" => \$host,
	"p|port=i" => \$port,
	"u|url=s",
	"v|verbose+" => \$verbose,
	"help+" => \$help,
	"h+" => \$shorthelp,
	"V|version+" => \$version) or &usage and exit(1);

    $help and &usage and exit(0);
    $shorthelp and &shortusage and exit(0);
    $version and &version and exit(0);
    $timeout=0 if ($timeout<0);
    $warning=0 if ($warning<0);
}

# Exit function: prints nagios status and dumps log
sub nagios_exit() {
    my $rc=nagstat->get_status();

    # Logging object
    logger->get_log();

    exit $rc;
}

# Signal handler for SIGALRM
sub alarm_handler() {
    my ($sig)=@_;
    logger->log_func(2,"Timeout exceeded\n");
    if ($stage>$stages{'presock'} && $stage<$stages{'sockclosed'})  {
	logger->log_func(2,"Socket has been opened, closing it\n");
	close ($sock);
    }
    nagstat->set_status(2,"probe timeout exceeded");
    nagios_exit;
}

# Signal handler for SIGINT and SIGTERM
sub int_handler()	{
    my ($sig)=@_;
    logger->log_func(2,"Caught signal ".$sig."\n");
    if ($stage>$stages{'presock'} && $stage<$stages{'sockclosed'})  {
	logger->log_func(2,"Socket has been opened, closing it\n");
	close ($sock);
    }
    nagstat->set_status(2,"caught signal ".$sig);
    nagios_exit;
}

# Actual probe opening socket to EES, sending message, receiving response and
# parsing the result
sub call_ees($)    {
    my $msg=shift;
    my $t1;
    my $t2;

    # Make sure to have starttime
    $t1=time();

    # Set alarm
    alarm($timeout);
    
    logger->log_func(2,"Opening connection to ".$host.":".$port."\n");
    # Open socket to $host:$port
    $sock = IO::Socket::INET->new(
	PeerAddr => $host,
	PeerPort => $port,
	Proto => 'tcp'
    );
    if (!defined $sock)	{
	nagstat->set_status(2,"Failed to connect ($!)");
	return 1;
    }
    $stage=$stages{'sockopen'};

    # Send soap message
    logger->log_func(3,"Socket opened, sending message\n");
    print $sock $msg;
    $stage=$stages{'datasent'};

    # Get header and response
    logger->log_func(3,"Message sent, waiting for response\n");
    my $header = <$sock>;
    if (!defined $header)    {
	my $summary="cannot read from socket ($!)";
	close($sock);
	nagstat->set_status(2,$summary);
	return 1;
    }
    $stage=$stages{'headrcvd'};
    
    # Chop of any carriage-return or line-feed from header
    $header =~ s/[\r\n]+$//;
    logger->log_func(3,"Header ".$header." received\n");
   
   # Dump remaining response in log
    logger->log_func(3,"Reading remaining response\n");
    while(my $line=<$sock>) {
	logger->log_func(3,"$line");
    }
    logger->log_func(3,"\n");
    $stage=$stages{'resprcvd'};

    # Close socket
    logger->log_func(3,"Response finished, closing socket\n");
    close($sock);
    $stage=$stages{'sockclosed'};
    logger->log_func(3,"Socket closed\n");
    # We are done with the socket, we have timing statistics
    $t2=time();

    # Set performance data
    my $dt=int(($t2-$t1)*1000+0.5)/1000;
    nagstat->set_perfdata("${dt}s;$warning;$timeout;0");

    # Check header
    if ("$header" eq "HTTP/1.1 200 OK")   {
	if ($dt<$warning)  {
	    nagstat->set_status(0,"Success");
	    return 0;
	}
	nagstat->set_status(1,"EES is slow in responding");
	return 0;
    }
    # There was a problem, chop of HTTP/1.1 and set status
    if ("$header" =~ /HTTP\/1.1 .*/)    {
	nagstat->set_status(1,"unexpected answer from host ($header)");
    } else {
	nagstat->set_status(2,"not a valid response ($header)");
    }
    return 1;
}

my $msg= <<EOF;
<?xml version="1.0" encoding="UTF-8"?>
<SOAP-ENV:Envelope xmlns:SOAP-ENV="http://schemas.xmlsoap.org/soap/envelope/"
xmlns:SOAP-ENC="http://schemas.xmlsoap.org/soap/encoding/"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xmlns:xsd="http://www.w3.org/2001/XMLSchema"
xmlns:dsig="http://www.w3.org/2000/09/xmldsig#"
xmlns:saml="urn:oasis:names:tc:SAML:2.0:assertion"
xmlns:XACMLcontext="urn:oasis:names:tc:xacml:2.0:context:schema:os"
xmlns:XACMLassertion="urn:oasis:names:tc:xacml:2.0:profile:saml2.0:v2:schema:assertion"
xmlns:XACMLpolicy="urn:oasis:names:tc:xacml:2.0:policy:schema:os"
xmlns:xenc="http://www.w3.org/2001/04/xmlenc#"
xmlns:XACMLService="http://www.globus.org/security/XACMLAuthorization/bindings"
xmlns:XACMLsamlp="urn:oasis:names:tc:xacml:2.0:profile:saml2.0:v2:schema:protocol"
xmlns:samlp="urn:oasis:names:tc:SAML:2.0:protocol">
<SOAP-ENV:Body>
<XACMLsamlp:XACMLAuthzDecisionQuery CombinePolicies="true" ReturnContext="true"
InputContextOnly="false" IssueInstant="2010-03-25T14:55:01Z" Version="2.0"
ID="ID-1804289383">
<saml:Issuer xsi:type="saml:NameIDType"
Format="urn:oasis:names:tc:SAML:1.1:nameid-format:X509SubjectName">NetCat</saml:Issuer>
<XACMLcontext:Request xsi:type="XACMLcontext:RequestType">
<XACMLcontext:Action xsi:type="XACMLcontext:ActionType">
</XACMLcontext:Action>
</XACMLcontext:Request>
</XACMLsamlp:XACMLAuthzDecisionQuery>
</SOAP-ENV:Body>
</SOAP-ENV:Envelope>
EOF

# Parse commandline options
getopts();

# Initialize logger and set loglevel
logger->new($verbose);

# Initialize nagios status logger
nagstat->new();

# Initialize signal handling
probeipc->new(\&alarm_handler,\&int_handler,$timeout);

# run actual EES probe
call_ees($msg);

# Dump nagios status, log and exit
nagios_exit();

