diff --git a/mk_livestatus/errors.py b/mk_livestatus/errors.py new file mode 100644 index 0000000..fb048e1 --- /dev/null +++ b/mk_livestatus/errors.py @@ -0,0 +1,6 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + + +class LivestatusError(Exception): + pass diff --git a/mk_livestatus/livestatus.py b/mk_livestatus/livestatus.py index 25ccdf9..a1f39eb 100644 --- a/mk_livestatus/livestatus.py +++ b/mk_livestatus/livestatus.py @@ -2,6 +2,7 @@ # -*- coding: utf-8 -*- from __future__ import unicode_literals +from .errors import LivestatusError import socket import json @@ -27,7 +28,7 @@ def call(self): __call__ = call def __str__(self): - request = 'GET %s' % (self._resource) + request = 'GET %s\nResponseHeader: fixed16' % (self._resource) if self._columns and any(self._columns): request += '\nColumns: %s' % (' '.join(self._columns)) if self._filters: @@ -44,27 +45,63 @@ def filter(self, filter_str): self._filters.append(filter_str) return self + def close(self): + pass + class Socket(object): - def __init__(self, peer): + def __init__(self, peer, encoding="latin-1"): self.peer = peer + self.socket = None + self.fd = None + self.encoding = encoding def __getattr__(self, name): return Query(self, name) + def cleanup(self): + + try: + self.fd.close() + del(self.fd) + except Exception: + pass + try: + self.socket.close() + del(self.socket) + except Exception: + pass + def call(self, request): try: if len(self.peer) == 2: - s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 1) + self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 3) + self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5) else: - s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) - s.connect(self.peer) - s.send(request) - s.shutdown(socket.SHUT_WR) - rawdata = s.makefile().read() + self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + self.socket.connect(self.peer) + self.socket.send(request) + self.socket.shutdown(socket.SHUT_WR) + self.fd = self.socket.makefile() + rawdata = self.fd.read() + self.cleanup() + except Exception as err: + self.cleanup() + raise LivestatusError("Failed to connect to Livestatus. Reason: %s" % (err)) + else: if not rawdata: - return [] - data = json.loads(rawdata) + raise LivestatusError("Livestatus service returned no data.") + data = self.validateHeader(rawdata) + data = json.loads(data, encoding=self.encoding) return [dict(zip(data[0], value)) for value in data[1:]] - finally: - s.close() + + def validateHeader(self, rawdata): + header = rawdata[0:16] + if header[0:3] == "200": + data = rawdata[16:] + return data + else: + raise LivestatusError("The Livestatus query contained an error. Reason: %s" % (rawdata[16:]))