naemon-query.py

#!/usr/bin/env python

import sys
import socket


class naemon_qh:
    __slots__ = ['query_handler', 'socket', 'read_size']
    read_size = 4096
    socket = False
    query_handler = False

    def __init__(self, query_handler):
        self.query_handler = query_handler

    def query(self, query):
        """Ask a raw query to naemon' query handler socket, return the raw
        response as a lazily generated sequence."""
        if not query:
            print "Missing query argument"
            return
        try:
            self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            self.socket.connect(self.query_handler)

            self.socket.send(query + '\0')
            while True:
                try:
                    out = self.socket.recv(self.read_size)
                    if not out:
                        break
                    yield out
                except KeyboardInterrupt:
                    print "Good bye."
                    break
            self.socket.close()
        except socket.error, e:
            print ("Couldn't connect to naemon socket %s: %s" % (
                    self.query_handler, str(e)
                )
            )

    def format(self, data, rowsep='\n', pairsep=';', kvsep='='):
        """Lazily format a response into a sequence of dicts"""
        def todict(row):
            return dict(x.split(kvsep) for x in row.split(pairsep))
        rest = ''
        for block in data:
            block = rest + block
            rows = block.split(rowsep)
            for row in rows[:-1]:
                yield todict(row)
            rest = rows[-1]
        if rest:
            yield todict(rest)

    def get(self, query, rowsep='\n', pairsep=';', kvsep='='):
        """Ask a query to naemon' query handler socket, and return an object
        representing a "standard" naemon qh response"""
        resp = self.query(query)
        return self.format(resp, rowsep, pairsep, kvsep)


if __name__ == '__main__':
    socket_path = '/var/lib/naemon/naemon.qh'
    query_args = []

    for arg in sys.argv[1:]:
        if arg.startswith('--socket='):
            socket_path = arg.split('=', 1)[1]
            continue
        query_args.append(arg)

    query = ' '.join(query_args)
    qh = naemon_qh(socket_path)
    for block in qh.query(query):
        sys.stdout.write(block)
        sys.stdout.flush()