ansible/apt_key

#!/usr/bin/python
# -*- coding: utf-8 -*-

# (c) 2012, Michael DeHaan <michael.dehaan@gmail.com>
# (c) 2012, Jayson Vantuyl <jayson@aggressive.ly>
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible 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 Ansible.  If not, see <http://www.gnu.org/licenses/>.

DOCUMENTATION = '''
---
module: apt_key
author: Jayson Vantuyl
version_added: 1.0
short_description: Add or remove an apt key
description:
    - Add or remove an I(apt) key, optionally downloading it
notes:
    - doesn't download the key unless it really needs it
    - as a sanity check, downloaded key id must match the one specified
    - best practice is to specify the key id and the url
options:
    id:
        required: false
        default: none
        description:
            - identifier of key
    url:
        required: false
        default: none
        description:
            - url to retrieve key from.
    state:
        required: false
        choices: [ absent, present ]
        default: present
        description:
            - used to specify if key is being added or revoked
examples:
    - code: "apt_key: url=https://ftp-master.debian.org/keys/archive-key-6.0.asc state=present"
      description: Add an Apt signing key, uses whichever key is at the URL
    - code: "apt_key: id=473041FA url=https://ftp-master.debian.org/keys/archive-key-6.0.asc state=present"
      description: Add an Apt signing key, will not download if present
    - code: "apt_key: url=https://ftp-master.debian.org/keys/archive-key-6.0.asc state=absent"
      description: Remove an Apt signing key, uses whichever key is at the URL
    - code: "apt_key: id=473041FA state=absent"
      description: Remove a Apt specific signing key
'''

from urllib2 import urlopen, URLError
from traceback import format_exc
from subprocess import Popen, PIPE, call
from re import compile as re_compile
from distutils.spawn import find_executable
from os import environ
from sys import exc_info

match_key = re_compile("^gpg:.*key ([0-9a-fA-F]+):.*$")

REQUIRED_EXECUTABLES=['gpg', 'grep', 'apt-key']


def find_missing_binaries():
    return [missing for missing in REQUIRED_EXECUTABLES if not find_executable(missing)]


def get_key_ids(key_data):
    p = Popen("gpg --list-only --import -", shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE)
    (stdo, stde) = p.communicate(key_data)

    if p.returncode > 0:
        raise Exception("error running GPG to retrieve keys")

    output = stdo + stde

    for line in output.split('\n'):
        match = match_key.match(line)
        if match:
            yield match.group(1)


def key_present(key_id):
    return call("apt-key list | 2>&1 grep -q %s" % key_id, shell=True) == 0


def download_key(url):
    if url is None:
        raise Exception("Needed URL but none specified")
    connection = urlopen(url)
    if connection is None:
        raise Exception("error connecting to download key from %r" % url)
    return connection.read()


def add_key(key):
    p = Popen("apt-key add -", shell=True, stdin=PIPE, stdout=PIPE, stderr=PIPE)
    (_, _) = p.communicate(key)

    return p.returncode == 0


def remove_key(key_id):
    return call('apt-key del %s' % key_id, shell=True) == 0


def return_values(tb=False):
    if tb:
        return {'exception': format_exc()}
    else:
        return {}


# use cues from the environment to mock out functions for testing
if 'ANSIBLE_TEST_APT_KEY' in environ:
    orig_download_key = download_key
    KEY_ADDED=0
    KEY_REMOVED=0
    KEY_DOWNLOADED=0


    def download_key(url):
        global KEY_DOWNLOADED
        KEY_DOWNLOADED += 1
        return orig_download_key(url)


    def find_missing_binaries():
        return []


    def add_key(key):
        global KEY_ADDED
        KEY_ADDED += 1
        return True


    def remove_key(key_id):
        global KEY_REMOVED
        KEY_REMOVED += 1
        return True


    def return_values(tb=False):
        extra = dict(
            added=KEY_ADDED,
            removed=KEY_REMOVED,
            downloaded=KEY_DOWNLOADED
        )
        if tb:
            extra['exception'] = format_exc()
        return extra


    if environ.get('ANSIBLE_TEST_APT_KEY') == 'none':
        def key_present(key_id):
            return False
    else:
        def key_present(key_id):
            return key_id == environ['ANSIBLE_TEST_APT_KEY']


def main():
    module = AnsibleModule(
        argument_spec=dict(
            id=dict(required=False, default=None),
            url=dict(required=False),
            state=dict(required=False, choices=['present', 'absent'], default='present')
        )
    )

    expected_key_id = module.params['id']
    url = module.params['url']
    state = module.params['state']
    changed = False

    missing = find_missing_binaries()

    if missing:
        module.fail_json(msg="can't find needed binaries to run", missing=missing,
            **return_values())

    if state == 'present':
        if expected_key_id and key_present(expected_key_id):
            # key is present, nothing to do
            pass
        else:
            # download key
            try:
                key = download_key(url)
                (key_id,) = tuple(get_key_ids(key))  # TODO: support multiple key ids?
            except Exception:
                module.fail_json(
                    msg="error getting key id from url",
                    **return_values(True)
                )

            # sanity check downloaded key
            if expected_key_id and key_id != expected_key_id:
                module.fail_json(
                    msg="expected key id %s, got key id %s" % (expected_key_id, key_id),
                    **return_values()
                )

            # actually add key
            if key_present(key_id):
                changed=False
            elif add_key(key):
                changed=True
            else:
                module.fail_json(
                    msg="failed to add key id %s" % key_id,
                    **return_values()
                )
    elif state == 'absent':
        # optionally download the key and get the id
        if not expected_key_id:
            try:
                key = download_key(url)
                (key_id,) = tuple(get_key_ids(key))  # TODO: support multiple key ids?
            except Exception:
                module.fail_json(
                    msg="error getting key id from url",
                    **return_values(True)
                )
        else:
            key_id = expected_key_id

        # actually remove key
        if key_present(key_id):
            if remove_key(key_id):
                changed=True
            else:
                module.fail_json(msg="error removing key_id", **return_values(True))
    else:
        module.fail_json(
            msg="unexpected state: %s" % state,
            **return_values()
        )

    module.exit_json(changed=changed, **return_values())

# include magic from lib/ansible/module_common.py
#<<INCLUDE_ANSIBLE_MODULE_COMMON>>
main()