Custom Ansible module to manage the entire VyOS config

example.yaml

---
- hosts: vyos
  vars:
    # This would ideally go in the host vars.
    vyos_config:
      # ...
      # The structure is identical to what `show configuration json pretty` in
      # operational mode produces. For the vast majority of VyOS commands, the
      # command <-> JSON translation is 1:1. For example, the command:
      #
      #  $ set system time-zone UTC
      #
      # has the following structural representation:
      #
      #   system:
      #     time-zone: UTC
      #
  tasks:
    - name: Set full system configuration
      vyos_full_config:
        config: '{{ vyos_config }}'

vyos_full_config.py

#!/usr/bin/python

# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU General Public License version 3 as published by the
# Free Software Foundation.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program. If not, see <https://www.gnu.org/licenses/>.

import json
import shlex

from ansible.module_utils.basic import AnsibleModule


TAG_BEGIN = '----- BEGIN -----'
TAG_END = '----- END -----'


def generate_diff(source, target, output, parents=[]):
    if source is not None and target is not None \
            and type(source) != type(target):
        raise ValueError('Source and target types do not match: %s != %s'
                         % (repr(source), repr(target)))

    if isinstance(source, list) or isinstance(target, list):
        source_items = set() if source is None else set(source)
        target_items = set() if target is None else set(target)

        for i in sorted(source_items - target_items):
            generate_diff(i, None, output, parents)
        for i in sorted(target_items - source_items):
            generate_diff(None, i, output, parents)
    elif isinstance(source, dict) or isinstance(target, dict):
        source_keys = set() if source is None else source.keys()
        target_keys = set() if target is None else target.keys()

        for k in sorted(source_keys - target_keys):
            # Remove the entire tree
            output.append(('delete', *parents, k))
        for k in sorted(source_keys & target_keys):
            generate_diff(source[k], target[k], output, parents + [k])
        for k in sorted(target_keys - source_keys):
            generate_diff(None, target[k], output, parents + [k])

        # Creating an empty dict
        if source is None and not target:
            output.append(('set', *parents))
    elif source != target:
        action = 'set' if target is not None else 'delete'
        value = target if target is not None else source
        output.append((action, *parents, value))


def extract_tagged(output):
    STATE_BEFORE = 0
    STATE_TAGGED = 1
    STATE_AFTER = 2

    before = []
    tagged = []
    after = []

    state = STATE_BEFORE

    for line in output.splitlines():
        if line == TAG_BEGIN:
            state = STATE_TAGGED
            continue
        elif line == TAG_END:
            if state != STATE_TAGGED:
                raise Exception('Encountered end tag before begin tag')
            state = STATE_AFTER
            continue

        if state == STATE_BEFORE:
            before.append(line)
        elif state == STATE_TAGGED:
            tagged.append(line)
        else:
            after.append(line)

    return '\n'.join(before), '\n'.join(tagged), '\n'.join(after)


def get_current_config(module, result):
    _, stdout, stderr = module.run_command(
        ['/opt/vyatta/bin/vyatta-op-cmd-wrapper', 'show', 'configuration', 'json'],
    )

    result['stdout'] = stdout
    result['stderr'] = stderr

    if stderr:
        raise Exception('Failed to query current configuration')

    return json.loads(stdout)


def run_config_commands(module, result, commands):
    script = [
        'source /opt/vyatta/etc/functions/script-template',
        'configure',
        *commands,
        shlex.join(('echo', TAG_BEGIN)),
        'compare',
        shlex.join(('echo', TAG_END)),
    ]

    if module.check_mode:
        script.append('discard')
    else:
        script.append('commit')
        script.append('save')
    script.append('exit')

    _, stdout, stderr = module.run_command(
        ['/bin/vbash', '-s'],
        data='\n'.join(script),
    )

    result['stdout'] = stdout
    result['stderr'] = stderr

    before, tagged, after = extract_tagged(stdout)

    if stderr or 'failed' in before or 'failed' in after:
        raise Exception('Failed to run configuration commands')
    elif not tagged:
        raise Exception('Expected changes, but none reported')

    if module._diff:
        result['diff'] = {
            'prepared': tagged,
        }


def main():
    module = AnsibleModule(
        argument_spec=dict(
            config=dict(type='dict', required=True),
        ),
        supports_check_mode=True
    )

    result = dict(changed=False)

    try:
        source_config = get_current_config(module, result)
        target_config = module.params['config']

        apply_commands = []
        generate_diff(source_config, target_config, apply_commands)

        if apply_commands:
            result['changed'] = True
            result['commands'] = [shlex.join(c) for c in apply_commands]

            run_config_commands(module, result, result['commands'])

        module.exit_json(**result)
    except Exception as e:
        module.fail_json(msg=str(e), **result)


if __name__ == '__main__':
    main()