build-zones

#!/usr/bin/env python3
"""Rebuild DNS config based on LDAP and template files.

Right now, it generates A/AAAA, CNAME, and reverse-DNS records for hosts, and
the result needs to be checked in to git.
"""
import time
from collections import namedtuple
from ipaddress import ip_address
from itertools import chain
from operator import attrgetter

import ldap3
from cached_property import cached_property
from ocflib.infra.ldap import ldap_ocf
from ocflib.infra.ldap import OCF_LDAP_HOSTS

BANNER = '''\
; @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
;                THIS FILE IS AUTOMATICALLY GENERATED
; @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
;
;   You should either make changes in LDAP and run `make`, or modify
;   the template under `templates`.
;
; @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
'''

RDNS_V4_FILE = 'etc/zones/db.226.229.169.in-addr.arpa'
RDNS_V6_FILE = 'etc/zones/db.0.0.0.0.1.0.8.8.0.4.1.f.7.0.6.2.ip6.arpa'

LE_ROOT = '_acme-challenge\tIN CNAME\tletsencrypt.ocf.io.'
LE_RECORD = '_acme-challenge.{record}\tIN CNAME\t{record}.letsencrypt.ocf.io.'


class Host(namedtuple('Host', ('hostname',))):
    def __new__(cls, attributes):
        return super().__new__(cls, attributes['cn'][0])

    def __init__(self, attributes):
        self.ldap_attrs = attributes

    @cached_property
    def description(self):
        d = '{self.hostname} ({self.ldap_attrs[type]})'.format(self=self)
        return d

    @cached_property
    def ipv4(self):
        ipv4, = self.ldap_attrs['ipHostNumber']
        ipv4 = ip_address(ipv4)
        return ipv4

    @cached_property
    def ipv6(self):
        if 'ip6HostNumber' in self.ldap_attrs:
            ipv6, = self.ldap_attrs['ip6HostNumber']
            ipv6 = ip_address(ipv6)
            # assert is_ocf_ip(ipv6)
            return ipv6

    @cached_property
    def last_quad(self):
        """Return the last quad from this host's IP address."""
        return self.ipv4.packed[3]

    def forward_dns_records(self):
        """Return forward DNS records."""

        def add_letsencrypt(record):
            if record == '@':
                return LE_ROOT.format(record=record)
            elif record.startswith('*.') or \
                    self.ldap_attrs['type'] in {'server', 'staffvm', 'vip'}:
                record = record.lstrip('*.')
                return LE_RECORD.format(record=record)
            else:
                return ''

        yield add_letsencrypt(self.hostname)
        if self.ipv4:
            yield '{self.hostname}\tIN A\t{self.ipv4}'.format(self=self)
        if self.ipv6:
            yield '{self.hostname}\tIN AAAA\t{self.ipv6}'.format(self=self)
        for cname in self.ldap_attrs.get('dnsCname', []):
            yield add_letsencrypt(cname)
            yield '{}\tIN CNAME\t{self.hostname}'.format(cname, self=self)
        for a in self.ldap_attrs.get('dnsA', []):
            # These are like CNAMEs but for special cases where CNAMEs are
            # illegal. Some examples: domain root, things used in MX or NS
            # records.
            yield add_letsencrypt(a)
            if self.ipv4:
                yield '{}\tIN A\t{self.ipv4}'.format(a, self=self)
            if self.ipv6:
                yield '{}\tIN AAAA\t{self.ipv6}'.format(a, self=self)

    def reverse_dns_records_v4(self):
        """Return reverse IPv4 DNS records."""
        if self.ipv4:
            yield '{self.last_quad} IN PTR ' \
                  '{self.hostname}.OCF.Berkeley.EDU.'.format(self=self)

    def reverse_dns_records_v6(self):
        """Return reverse IPv6 DNS records."""
        last_quad_dotted = '.'.join(reversed(str(self.last_quad).zfill(4)))
        # TODO: uppercase these properly like IPv4
        if self.ipv6:
            yield '{} IN PTR {self.hostname}.ocf.berkeley.edu.'.format(
                last_quad_dotted,
                self=self,
            )

    def __str__(self):
        return '{self.hostname}\tIN A\t{self.ipv4}'.format(self=self)


def hosts_from_ldap():
    """Return set of hosts at the OCF."""
    with ldap_ocf() as c:
        c.search(
            OCF_LDAP_HOSTS,
            '(cn=*)',
            attributes=ldap3.ALL_ATTRIBUTES,
        )
        hosts = {
            Host(record['attributes'])
            for record in c.response
        }
    return hosts


def forward_file(hosts, serial):
    return open('templates/db.ocf.tmpl').read().format(
        banner=BANNER,
        serial=serial,
        records='\n\n'.join(
            '\n'.join(
                ['; {}'.format(host.description)] +
                list(filter(None, host.forward_dns_records())),
            )
            for host in sorted(hosts, key=attrgetter('ipv4'))
        ),
    )


def reverse_file_v4(hosts, serial):
    return open('templates/db.ocf.in-addr.arpa.tmpl').read().format(
        banner=BANNER,
        serial=serial,
        records='\n'.join(
            chain.from_iterable(
                host.reverse_dns_records_v4()
                for host in sorted(hosts, key=attrgetter('ipv4'))
            ),
        ),
    )


def reverse_file_v6(hosts, serial):
    return open('templates/db.ocf.ip6.arpa.tmpl').read().format(
        banner=BANNER,
        serial=serial,
        records='\n'.join(
            chain.from_iterable(
                host.reverse_dns_records_v6()
                for host in sorted(
                    filter(lambda host: host.ipv6, hosts),
                    key=attrgetter('ipv6'),
                )
            ),
        ),
    )


def next_serial():
    """Return the next-greatest serial number from what's currently in use.

    We don't want to assume networking is working, or that the DNS server is
    up, or that the most recent changes have already propagated (e.g. they
    might be half-way through a Jenkins pipeline).

    So, instead we read it from a file we check into git.
    """
    date_as_str = time.strftime('%Y%m%d')
    try:
        serial = open('etc/serial').read().strip()
        if serial.startswith(date_as_str) and len(serial) == 10:
            return str(int(serial) + 1)
    except IOError:
        pass
    return date_as_str + '00'


def main(argv=None):
    hosts = hosts_from_ldap()

    # generate zone files
    serial = next_serial()

    with open('etc/serial', 'w') as f:
        f.write('{}\n'.format(serial))

    ocf_zone = forward_file(hosts, serial)
    with open('etc/db.ocf', 'w') as f:
        f.write(ocf_zone)

    reverse_v4_zone = reverse_file_v4(hosts, serial)
    with open(RDNS_V4_FILE, 'w') as f:
        f.write(reverse_v4_zone)

    reverse_v6_zone = reverse_file_v6(hosts, serial)
    with open(RDNS_V6_FILE, 'w') as f:
        f.write(reverse_v6_zone)


if __name__ == '__main__':
    exit(main())