mon monitor plugin to test msgpack-rpc services

pymsgpack-rpc.monitor.py

#!/usr/bin/env python
#
# pymsgpack-rpc.monitor - mon monitor plugin to test msgpack-rpc services
#

from __future__ import with_statement
import logging
import msgpack
import optparse
import os
import random
import socket
import sys
import time

try:
    import json
except ImportError:
    import simplejson as json

REQUEST = 0
RESPONSE = 1
NOTIFY = 2

def connect(host, port, timeout=10):
    res = socket.getaddrinfo(host, port, socket.AF_UNSPEC,
                             socket.SOCK_STREAM, 0, socket.AI_PASSIVE)[0]
    af, socktype, proto, canonname, sockaddr = res
    sock = socket.socket(af, socktype, proto)
    if af == socket.AF_INET6:
        if hasattr(socket, "IPPROTO_IPV6"):
            sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
    sock.settimeout(timeout)
    sock.connect((host, port))
    return sock

def call(host, port, message, args=None, timeout=10):
    sock = None
    try:
        sock = connect(host, port, timeout=timeout)
#       msgid = int(random.random() * 0xFFFFFFFF)
        msgid = int(time.time()) & 0xFFFFFFFF
        req = msgpack.packb(([REQUEST, msgid, message, args]))
        assert len(req) == sock.send(req)
        _message = msgpack.unpackb(sock.recv(4096))
        _msgsize = len(_message)
        if _msgsize != 4 and _msgsize != 3:
            raise RuntimeError("Invalid MessagePack-RPC protocol: message = {0}".format(_message))
        _msgtype = _message[0]
        if _msgtype == RESPONSE:
            _msgid, _error, _result = _message[1], _message[2], _message[3]
            if _msgid == msgid:
                if _error:
                    raise RuntimeError("Error: %s" % (_error,))
                else:
                    return _result
            else:
                raise RuntimeError("Invalid message id: %d != %d" % (msgid, _msgid))
        else:
            raise RuntimeError("Invalid message type: %d != %d" % (_msgtype, RESPONSE))
    finally:
        if sock:
            sock.close()

def main(args):
    parser = optparse.OptionParser("usage %prog [OPTIONS]", add_help_option=False)
    parser.add_option('-a', '--argument', default=[], action='append', dest='arguments', help='msgpack-rpc argument, in JSON format.')
    parser.add_option('-e', '--expects', default='ok', dest='expects', help='expected return value from msgpkac-rpc.')
    parser.add_option('-m', '--message', default='ping', dest='message', help='msgpack-rpc message.')
    parser.add_option('-p', '--port', type='int', default=[], action='append', dest='ports', help='msgpack-rpc port.')
    parser.add_option('-t', '--timeout', type='int', default=10, dest='timeout', help='timeout')
    parser.add_option('--help', action='help', help='show this message and exit.')
    parser.add_option('-v', '--verbose', default=False, action='store_true', dest='verbose', help='show verbose message.')
    (options, args) = parser.parse_args(args)

    if options.verbose:
        logging.basicConfig(level=logging.DEBUG)

    arguments = [ json.loads(s) for s in options.arguments ]
    ports = options.ports

    f = 0
    for host, port in [ (host, port) for host in args[1:] for port in ports ]:
        try:
            got = call(host, port, options.message, args=arguments, timeout=options.timeout)
            if got == options.expects:
                logging.debug("Got response from %s:%d#%s ==> %s" % (host, port, options.message, repr(got)))
            else:
                logging.error("Invalid response from %s:%d#%s (got %s, expected %s)" % (host, port, options.message, repr(got), repr(options.expects)))
                f += 1
        except Exception, error:
            logging.error("Error during calling %s:%d#%s: %s: %s" % (host, port, options.message, error.__class__.__name__, error.message))
            f += 1

    sys.exit(0 if f == 0 else 1)

if __name__ == "__main__":
    main(sys.argv)

# vim:set ft=python :