status-bot/modules/batcave/client.py

import json
import logging

from urllib2 import urlopen, URLError


class BatcaveClient(object):
    def __init__(self, url):
        assert url.startswith("http://") or url.startswith("https://"), \
            "BATCAVE URL must use http(s) protocol"
        assert url.endswith("/"), "BATCAVE URL must end with slash."

        self.base_url = url
        self.logger = logging.getLogger('BATCAVE')

    def __load_response(self, url, error_context):
        raw_data = None
        try:
            raw_data = urlopen(self.base_url + url)
        except URLError as err:
            self.logger.error("Failed to contact BATCAVE for %s: %s",
                              error_context, err)
            return None

        try:
            return json.load(raw_data)
        except ValueError as err:
            self.logger.error("Could not parse response for %s: %s",
                              error_context, err)
            return None

    def get_nodes(self):
        url = 'nodes.json'
        return self.__load_response(url, 'nodes')

    def get_node(self, nodeid):
        """Query the given node's data from the BATCAVE."""

        url = 'node/{0}.json'.format(nodeid)
        return self.__load_response(url, 'node \'' + nodeid + '\'')

    def find_node_by_name(self, name, fuzzymatch=True, single_match_only=True):
        """Tries to find a node by given name."""

        url = 'find?name=' + name + '&fuzzy=' + ('1' if fuzzymatch else '0')
        matches = self.__load_response(url, 'find_name=' + name)

        if single_match_only:
            if len(matches) == 1:
                return matches[0]
            else:
                return None

        return matches

    def find_node_by_mac(self, mac, single_match_only=True):
        """Tries to find a node by given MAC address."""

        url = 'find?mac=' + mac
        matches = self.__load_response(url, 'find_mac=' + mac)

        if single_match_only:
            if len(matches) == 1:
                return matches[0]
            else:
                return None

        return matches

    def get_nodefield(self, nodeid, field):
        """Query the given field for the given nodeid from the BATCAVE."""

        url = self.base_url + 'node/{0}/{1}'.format(nodeid, field)
        ctx = "node '{0}'->'{1}'".format(nodeid, field)

        return self.__load_response(url, ctx)

    def get_providers(self):
        url = 'providers?format=json'
        return self.__load_response(url, 'providers')

    def get_status(self):
        url = 'status'
        return self.__load_response(url, 'status')