#!/usr/bin/python
# -*- coding:utf-8 -*-

###############################################################################
# filename  : check_wowza_edge.py
# contact   : emrah.com@gmail.com
#
# This is a Nagios plugin. It checks the Wowza edge server status via the Wowza
# load balancer server.
#
# Dependencies:
#   python-lxml python-ipaddr
#
# Usage:
#   check_wowza_edge.py -h
#
# Nagios usage:
#   command_line /path/to/check_wowza_edge.py -H '$HOSTADDRESS$' -e '$ARG1$'
#                                             -w '$ARG2$' -c '$ARG3$'
###############################################################################

import sys
import urllib2
import ipaddr
import argparse
from lxml import etree

# Load balancer's XML info page.
STATUS_PAGE = '/loadbalancer?serverInfoXML'
# Debug mode.
DEBUG = False



#------------------------------------------------------------------------------
def argument_parser():
    '''Parse the parameters.'''

    try:
        # Parse parameters.
        parser = argparse.ArgumentParser()
        parser.add_argument('-H', '--hostname', required=True,
               help='Load balancer\'s host name or IP Address')
        parser.add_argument('-p', '--port', type=int, default=1935,
               help='Load balancer\'s port number (default: 1935)')
        parser.add_argument('-e', '--edge', required=True,
               help='Edge\'s host name or IP Address')
        parser.add_argument('-w', '--warning', type=int, default=500,
               help='The connection count in warning status (default: 500)')
        parser.add_argument('-c', '--critical', type=int, default=750,
               help='The connection count in critical status (default: 750)')
        parser.add_argument('-t', '--timeout', type=int, default=10,
               help='The connection timeout (default: 10)')

        result = parser.parse_args()
    except Exception:
        if DEBUG: raise

        result = None
    return result



#------------------------------------------------------------------------------
def argument_validator(args):
    '''Validate the parameters.'''

    try:
        # Validate parameters.
        if (not ipaddr.IPAddress(args.hostname)):
            raise NameError('Invalid IP address for the load balancer')
        elif (not ipaddr.IPAddress(args.edge)):
            raise NameError('Invalid IP address for the edge')
        elif (not args.port > 0 or not args.port < 65536):
            raise NameError('Invalid port number for the load balancer')
        elif (not args.warning > 0):
            raise NameError('Invalid connection count for the warning status')
        elif (not args.critical > 0):
            raise NameError('Invalid connection count for the critical status')
        elif (not (args.warning < args.critical)):
            raise NameError('The connection count for the warning status ' +
                            'must be smaller than the connection count for ' +
                            'the critical status')
        elif (not args.timeout > 0):
            raise NameError('Invalid connection timeout')

        result = args
    except Exception:
        if DEBUG: raise

        result = None
    return result



# -----------------------------------------------------------------------------
def get_edge_info(loadbalancer_ip, loadbalancer_port, edge_ip, timeout):
    ''' Get the Wowza edge server info via the Wowza load balancer. '''

    try:
        # Get XML info page's content.
        url = 'http://%s:%s/%s' % (loadbalancer_ip, loadbalancer_port,
                                   STATUS_PAGE)
        req = urllib2.urlopen(url, timeout=timeout)
        res = etree.parse(req)
        req.close()

        # Find node for the given edge server.
        for app in res.xpath('LoadBalancerServer'):
            if (app.findtext('redirect') != edge_ip):
                continue

            # Get the info for the given edge server.
            edge = {}
            edge['status'] = app.findtext('status')
            edge['connect_count'] = long(app.findtext('connectCount'))
            edge['last_message'] = app.findtext('lastMessage')
            edge['server_id'] = app.findtext('serverId')
            edge['redirect'] = app.findtext('redirect')
            edge['redirect_count'] = long(app.findtext('redirectCount'))
            break

        result = edge
    except Exception:
        if DEBUG: raise

        result = None

    return result



# -----------------------------------------------------------------------------
def get_nagios_status(args, edge):
    '''
    Get the Nagios plugin's return variables.
    It returns (code, message) pair.
    '''

    try:
        if edge['status'] == 'RUNNING':
            if (edge['connect_count'] > args.critical):
                code = 2
                status = 'CRITICAL'
            elif (edge['connect_count'] > args.warning):
                code = 1
                status = 'WARNING'
            else:
                code = 0
                status = 'OK'
        else:
            code = 2
            status = 'CRITICAL'

        message = '%s - %s, %s connections, %s redirects, %s ago.\n' \
                % (status, edge['status'], edge['connect_count'],
                   edge['redirect_count'], edge['last_message'])

        result = {'code'    : code,
                  'message' : message}
    except Exception:
        if DEBUG: raise

        result = None

    return result



# -----------------------------------------------------------------------------
# Program main.
# -----------------------------------------------------------------------------
if __name__ == '__main__':
    try:
        # Get parameters.
        args = argument_parser()
        if args is None:
            raise NameError('Invalid parameters')

        # Validate parameters.
        args = argument_validator(args)
        if args is None:
            raise NameError('Could not validate the parameters')

        # Get the edge server info.
        edge = get_edge_info(args.hostname, args.port, args.edge, args.timeout)
        if edge is None:
            raise NameError('Could not get the edge server info')

        # Get the Nagios plugin's status for this edge server.
        nagios = get_nagios_status(args, edge)
        if nagios is None:
            raise NameError('Could not get the Nagios plugin\'s status')
    except Exception, err:
        if DEBUG: raise

        nagios = {'code'    : 3,
                  'message' : 'UNKNOWN - %s\n' % (str(err))}

    sys.stdout.write(nagios.get('message', 'UNKNOWN'))
    sys.exit(nagios.get('code', 3))