eid-ops/scripts/cosmosdata/datasource.py

import os
import re
import sys
import yaml
import socket


class DataSource(object):

    def __init__(self):
        self._groups = {}

    def lookup(self, hostname):
        raise NotImplemented('subclass must implement lookup')

    def all_hosts(self):
        raise NotImplemented('subclass must implement all_hosts')

    def groups(self):
        """
        Return {group: {members: [hosts]}} data that is used in creating hostgroups
        """
        return self._groups

    def add_members_to_group(self, group, hosts):
        """
        Add hosts (FQDNs) to a group.
        :type group: str
        :type hosts: [str]
        """
        if not hosts:
            return
        if group not in self._groups:
            self._groups[group] = {'members': []}
        for this in hosts:
            if this not in self._groups[group]['members']:
                self._groups[group]['members'] += [this]

    def load_classes(self, filename):
        with open(filename) as fd:
            _classes = yaml.safe_load(fd)
            for this_r, classes in _classes.items():
                r_comp = re.compile(this_r)
                hosts = []
                for host in self.all_hosts():
                    if re.match(r_comp, host):
                        hosts += [host]
                if not classes:
                    continue
                for cls in classes.keys():
                    parts = cls.split('::')
                    if parts[0] in ['sunet', 'eduid', 'thulin']:
                        parts = parts[1:]
                    cls = '_'.join(parts)
                    self.add_members_to_group(cls, hosts)


class DNSResolver(DataSource):
    """
    DNS based hostname-to-IPs resolver, returning both IPv4 and IPv6 addresses.
    """

    def __init__(self, suffixes):
        super(DNSResolver, self).__init__()
        self._suffixes = suffixes

    def lookup(self, hostname):
        res = gai4 = gai6 = []
        try:
            gai4 = socket.getaddrinfo(hostname, 0, socket.AF_INET, 0, socket.IPPROTO_TCP)
        except socket.gaierror:
            sys.stderr.write('Host {} not found in DNS (v4)\n'.format(hostname))
        try:
            gai6 = socket.getaddrinfo(hostname, 0, socket.AF_INET6, 0, socket.IPPROTO_TCP)
        except socket.gaierror:
            sys.stderr.write('Host {} not found in DNS (v6)\n'.format(hostname))
        for this in gai4 + gai6:
            res.append(this[4][0])
        return res

    def all_hosts(self):
        """
        Return a list of all hosts in this cosmos environment.

        Currently this is done by looking for subdirectorys with a '.eduid.se' in them,
        but not starting with a '.'.
        """
        hosts = filter(lambda fn: os.path.isdir(fn) and \
                       not fn.startswith('.'),

                       os.listdir('.'))
        matching = []
        for h in hosts:
            for s in self._suffixes:
                if h.endswith(s):
                    matching += [h]
                    break
        return sorted(matching)


class HostsResolver(DataSource):
    """
    Return data from a YAML file with entrys like this:

    site_hosts:
       host1_vlan10:
          addrs:
          - 192.168.10.1
          dhcp: false
          fqdn: host1.example.org
          mac: xxx
       host2_vlan2:
          addrs:
            - 192.168.2.2
            - 2001::1
            - 2001::2
          dhcp: true
          fqdn: host2.example.org
          mac: yyy
          type: server
    """

    def __init__(self, fn='global/overlay/etc/puppet/hiera/data/hosts.yaml'):
        super(HostsResolver, self).__init__()
        self._fn = fn
        self._data = _load_hosts_file(fn)

    def lookup(self, hostname):
        hostname = hostname.lower()
        res = []
        for this, values in self._data.items():
            fqdn = self._get_fqdn(this)
            if not fqdn:
                continue
            if fqdn == hostname:
                res += values.get('addrs', [])
        return res

    def all_hosts(self):
        """
        Return a list of all FQDNs in this cosmos environment.
        """
        res = {}
        for this in self._data.keys():
            fqdn = self._get_fqdn(this)
            if not fqdn:
                continue
            res[fqdn] = True
        return res.keys()

    def groups(self):
        """
        Create hostgroups based on 'type' in the hosts data.
        """
        res = super(HostsResolver, self).groups()
        for this, values in self._data.items():
            if 'type' in values:
                group = values['type'].lower()
                fqdn = self._get_fqdn(this)
                if fqdn is None:
                    continue
                self.add_members_to_group(group, [fqdn])
        return res

    def _get_fqdn(self, key):
        fqdn = self._data[key]['fqdn'].lower()
        if '-old.' in fqdn:
            return None
        if self._data[key].get('monitor') is False:
            return None
        return fqdn


def _load_hosts_file(fn):
    """
    Load site-hosts data from hosts.yaml.
    """
    with open(fn) as fd:
        data = yaml.load(fd)
    return data['site_hosts']