aws_weighted_dns_by_id.py

"""Handle generic DNS events."""
import boto3
import botocore.exceptions
import logging

logger = logging.getLogger()


class DNS(object):
    """Create or delete DNS records in Route53."""

    def __init__(self, hosted_zone_id, record_ttl=None, record_weight=None):
        """Create a handler which manages records in the provided zone. Records will be created
        with TTL set to ``record_ttl`` seconds and a weight of ``record_weight``.
        """
        self.conn = boto3.client('route53')
        self.hosted_zone_id = hosted_zone_id
        self.hosted_zone_name = get_hosted_zone_name(self.conn, self.hosted_zone_id)
        self.record_ttl = record_ttl or 30
        self.record_weight = record_weight or 10

    def list_records(self, identifier):
        """List Route53 records which have the given ``identifier``."""
        records = []
        next_record_identifier = None
        while True:
            req = dict(
                HostedZoneId=self.hosted_zone_id,
            )
            if next_record_identifier:
                req.update(StartRecordIdentifier=next_record_identifier)

            res = self.conn.list_resource_record_sets(**req)
            for record in res['ResourceRecordSets']:
                if record.get('SetIdentifier') == identifier:
                    records.append(record)

            next_record_identifier = res.get('NextRecordIdentifier')
            if not next_record_identifier:
                break
        return records

    def create(self, identifier, domain, address):
        """Create a DNS record."""
        res = self.conn.change_resource_record_sets(
            HostedZoneId=self.hosted_zone_id,
            ChangeBatch=dict(
                Comment="add record for {}".format(identifier),
                Changes=[
                    dict(
                        Action='UPSERT',
                        ResourceRecordSet=dict(
                            Name=domain,
                            Type='A',
                            TTL=self.record_ttl,
                            SetIdentifier=identifier,
                            Weight=self.record_weight,
                            ResourceRecords=[
                                dict(Value=address),
                            ],
                        ),
                    ),
                ],
            ),
        )
        change_id = res['ChangeInfo']['Id']
        msg = "created record identifier={} domain={} address={} change_id={}"
        logger.info(msg.format(identifier, domain, address, change_id))

    def delete(self, identifier):
        """Delete DNS records."""
        changes = []
        for record in self.list_records(identifier):
            changes.append(dict(
                Action='DELETE',
                ResourceRecordSet=record,
            ))

        if changes:
            conn = boto3.client('route53')
            res = conn.change_resource_record_sets(
                HostedZoneId=self.hosted_zone_id,
                ChangeBatch=dict(
                    Comment="delete records for {}".format(identifier),
                    Changes=changes,
                ),
            )
            change_id = res['ChangeInfo']['Id']
            msg = "deleted records identifier={} change_id={}"
            logger.info(msg.format(identifier, change_id))
        else:
            msg = "no records to delete identifier={}"
            logger.info(msg.format(identifier))


def get_hosted_zone_name(conn, hosted_zone_id):
    """Return the domain name for the hosted zone."""
    try:
        res = conn.get_hosted_zone(Id=hosted_zone_id)
    except botocore.exceptions.ClientError as e:
        if e.response['ResponseMetadata']['HTTPStatusCode'] == 404:
            return None
        raise e
    return res['HostedZone']['Name']