digital_ocean_dns_record.py

#!/usr/bin/python
# vim: fileencoding=utf-8 tw=100 expandtab ts=4 sw=4 :

from __future__ import unicode_literals


from unittest import TestCase

try:
    from dopy.manager import DoError, DoManager
    HAS_DOPY = True
except ImportError:
    DoError = object()
    DoManager = object()
    HAS_DOPY = False


TEST_ENTRIES = [
    {
        'name': '@',
        'type': 'NS',
        'id': 1,
        'data': 'ns1.digitalocean.com',
    },
    {
        'name': '@',
        'type': 'NS',
        'id': 2,
        'data': 'ns2.digitalocean.com',
    },
    {
        'name': '@',
        'type': 'NS',
        'id': 3,
        'data': 'ns3.digitalocean.com',
    },
    {
        'name': '@',
        'type': 'A',
        'id': 4,
        'data': '1.1.1.1',
    },
    {
        'name': '@',
        'type': 'A',
        'id': 5,
        'data': '2.2.2.2',
    },
    {
        'name': '*',
        'type': 'CNAME',
        'id': 6,
        'data': '@',
    }
]


def guess_api_token(module):
    return module.params['api_token'] or os.environ['DO_API_TOKEN'] or os.environ['DO_API_KEY']


class DoDnsRecordSync(object):
    def __init__(self, manager, domain):
        self.manager = manager
        self.domain = domain

    def filter_entries(self, entries, name, type_, priority=None, port=None, weight=None):
        for entry in entries:
            if entry['name'] == name \
                    and entry['type'] == type_ \
                    and (priority is None or entry['priority'] == priority) \
                    and (port is None or entry['port'] == port) \
                    and (weight is None or entry['weight'] == weight):
                yield entry

    def ensure(self, name, values, type_, priority=None, port=None, weight=None):
        entries = list(
            self.filter_entries(
                self.manager.all_domain_records(self.domain),
                name,
                type_,
                priority,
                port,
                weight,
            )
        )

        entries_idx = {x['data']: x for x in entries}

        present = set(entries_idx.keys())
        expected = set(values)

        for value in (present - expected):
            self.manager.destroy_domain_record(self.domain, entries_idx[value]['id'])

        for value in (expected - present):
            kwargs = {
                'domain_id': self.domain,
                'record_type': type_,
                'data': value,
                'name': name,
            }

            if priority is not None:
                kwargs['priority'] = priority

            if port is not None:
                kwargs['port'] = port

            if weight is not None:
                kwargs['weight'] = weight

            self.manager.new_domain_record(**kwargs)

        return present != expected


class DoDnsRecordSyncTest(TestCase):
    def test_present_add(self):
        import mock

        m = mock.Mock()
        m.all_domain_records = mock.MagicMock(return_value=TEST_ENTRIES)
        m.new_domain_record = mock.MagicMock()

        s = DoDnsRecordSync(m, 'test.domain')
        c = s.ensure(
            '@',
            ['1.1.1.1', '2.2.2.2', '3.3.3.3'],
            'A'
        )

        m.new_domain_record.assert_called_once_with(
            domain_id='test.domain',
            record_type='A',
            data='3.3.3.3',
            name='@'
        )
        self.assertTrue(c)

    def test_present_remove(self):
        import mock

        m = mock.Mock()
        m.all_domain_records = mock.MagicMock(return_value=TEST_ENTRIES)
        m.destroy_domain_record = mock.MagicMock()

        s = DoDnsRecordSync(m, 'test.domain')
        c = s.ensure(
            '@',
            ['1.1.1.1'],
            'A'
        )

        m.destroy_domain_record.assert_called_once_with('test.domain', 5)
        self.assertTrue(c)

    def test_do_nothing(self):
        import mock

        m = mock.Mock()
        m.all_domain_records = mock.MagicMock(return_value=TEST_ENTRIES)

        s = DoDnsRecordSync(m, 'test.domain')
        c = s.ensure(
            '@',
            ['1.1.1.1', '2.2.2.2'],
            'A'
        )

        self.assertFalse(c)


def main():
    module = AnsibleModule(
        argument_spec={
            'state': {
                'default': 'present',
                'choices': ['present'],
            },
            'domain': {
                'required': True,
            },
            'name': {
                'required': True,
            },
            'values': {
                'type': 'list',
                'default': [],
            },
            'type': {
                'required': True,
            },
            'api_token': {
                'aliases': ['API_TOKEN'],
                'no_log': True,
            },
            'priority': {
                'type': 'int',
                'default': -1,
            },
            'port': {
                'type': 'int',
                'default': -1,
            },
            'weight': {
                'type': 'int',
                'default': -1,
            }
        },
    )

    if not HAS_DOPY:
        return module.fail_json(msg='dopy required for this module')

    m = DoManager(None, os.environ['DO_API_TOKEN'], 2)
    s = DoDnsRecordSync(m, module.params['domain'])

    priority = module.params['priority'] if module.params['priority'] >= 0 else None
    port = module.params['port'] if module.params['port'] >= 0 else None
    weight = module.params['weight'] if module.params['weight'] >= 0 else None

    return module.exit_json(
        changed=s.ensure(
            module.params['name'],
            module.params['values'],
            module.params['type'],
            priority,
            port,
            weight,
        ),
        msg='DNS records are up to date',
    )


from ansible.module_utils.basic import *
if __name__ == '__main__':
    main()