ansible/cloud/rax_files_objects

#!/usr/bin/python -tt

# (c) 2013, Paul Durivage <paul.durivage@rackspace.com>
#
# 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:
short_description: create, fetch, and delete objects in Rackspace Cloud Files
description:
  - Upload, download, and delete objects in Rackspace Cloud Files
version_added: "1.5"
options:
requirements: [ "pyrax" ]
author: Paul Durivage
notes:
  - Something
'''

EXAMPLES = ''' '''

import os

try:
    import pyrax
except ImportError, e:
    print("failed=True msg='pyrax is required for this module'")
    sys.exit(1)

EXIT_DICT = {'success': False}
META_PREFIX = 'x-object-meta-'


def _get_container(module, cf, container):
    try:
        return cf.get_container(container)
    except pyrax.exc.NoSuchContainer, e:
        module.fail_json(msg=e.message)


def upload(module, cf, container, src, dest, meta, expires):
    """ Uploads a single object or a folder to Cloud Files Optionally sets an
    metadata, TTL value (expires), or Content-Disposition and Content-Encoding
    headers.
    """
    c = _get_container(module, cf, container)

    num_objs_before = len(c.get_object_names())

    if not src:
        module.fail_json(msg='src must be specified when uploading')

    src = os.path.abspath(os.path.expanduser(src))
    is_dir = os.path.isdir(src)

    if not is_dir and not os.path.isfile(src) or not os.path.exists(src):
        module.fail_json(msg='src must be a file or a directory')
    if dest and is_dir:
        module.fail_json(msg='dest cannot be set when whole '
                             'directories are uploaded')

    cont_obj = None
    if dest and not is_dir:
        try:
            cont_obj = c.upload_file(src, obj_name=dest, ttl=expires)
        except Exception, e:
            module.fail_json(msg=e.message)
    elif is_dir:
        try:
            id, total_bytes = cf.upload_folder(src, container=c.name, ttl=expires)
        except Exception, e:
            module.fail_json(msg=e.message)

        while True:
            bytes = cf.get_uploaded(id)
            if bytes == total_bytes:
                break
            time.sleep(1)
    else:
        try:
            cont_obj = c.upload_file(src, ttl=expires)
        except Exception, e:
            module.fail_json(msg=e.message)

    num_objs_after = len(c.get_object_names())

    if not meta:
        meta = {}

    meta_result = {}
    if meta:
        if cont_obj:
            meta_result = cont_obj.set_metadata(meta)
        else:
            def _set_meta(objs, meta):
                """ Sets metadata on a list of objects specified by name """
                for obj in objs:
                    try:
                        result = c.get_object(obj).set_metadata(meta)
                    except Exception, e:
                        module.fail_json(msg=e.message)
                    else:
                        meta_result[obj] = result
                return meta_result

            def _walker(objs, path, filenames):
                """ Callback func for os.path.walk  """
                prefix = ''
                if path != src:
                    prefix = path.split(src)[-1].lstrip('/')
                filenames = [os.path.join(prefix, name) for name in filenames
                             if not os.path.isdir(name)]
                objs += filenames

            _objs = []
            os.path.walk(src, _walker, _objs)
            meta_result = _set_meta(_objs, meta)

    EXIT_DICT['success'] = True
    EXIT_DICT['container'] = c.name
    EXIT_DICT['msg'] = "Uploaded %s to container: %s" % (src, c.name)
    if cont_obj or locals().get('bytes'):
        EXIT_DICT['changed'] = True
    if meta_result:
        EXIT_DICT['meta'] = {
            'updated': True
        }

    if cont_obj:
        EXIT_DICT['bytes'] = cont_obj.total_bytes
        EXIT_DICT['etag'] = cont_obj.etag
    else:
        EXIT_DICT['bytes'] = total_bytes

    module.exit_json(**EXIT_DICT)


def download(module, cf, container, src, dest, structure):
    """ Download objects from Cloud Files to a local path specified by "dest".
    Optionally disable maintaining a directory structure by by passing a
    false value to "structure".
    """
    # Looking for an explicit destination
    if not dest:
        module.fail_json(msg='dest is a required argument when '
                             'downloading from Cloud Files')

    # Attempt to fetch the container by name
    c = _get_container(module, cf, container)

    # Accept a single object name or a comma-separated list of objs
    # If not specified, get the entire container
    if src:
        objs = src.split(',')
        objs = map(str.strip, objs)
    else:
        objs = c.get_object_names()

    dest = os.path.abspath(os.path.expanduser(dest))
    is_dir = os.path.isdir(dest)

    if not is_dir:
        module.fail_json(msg='dest must be a directory')

    results = []
    for obj in objs:
        try:
            c.download_object(obj, dest, structure=structure)
        except Exception, e:
            module.fail_json(msg=e.message)
        else:
            results.append(obj)

    len_results = len(results)
    len_objs = len(objs)

    EXIT_DICT['container'] = c.name
    EXIT_DICT['requested_downloaded'] = results
    if results:
        EXIT_DICT['changed'] = True
    if len_results == len_objs:
        EXIT_DICT['success'] = True
        EXIT_DICT['msg'] = "%s objects downloaded to %s" % (len_results, dest)
    else:
        EXIT_DICT['msg'] = "Error: only %s of %s objects were " \
                           "downloaded" % (len_results, len_objs)
    module.exit_json(**EXIT_DICT)


def delete(module, cf, container, src, dest):
    """ Delete specific objects by proving a single file name or a
    comma-separated list to src OR dest (but not both).  Ommitting file name(s)
    assumes the entire container is to be deleted.
    """
    objs = None
    if src and dest:
        module.fail_json(msg="Error: ambiguous instructions; files to be deleted "
                             "have been specified on both src and dest args")
    elif dest:
        objs = dest
    else:
        objs = src

    c = _get_container(module, cf, container)

    if objs:
        objs = objs.split(',')
        objs = map(str.strip, objs)
    else:
        objs = c.get_object_names()

    num_objs = len(objs)

    results = []
    for obj in objs:
        try:
            result = c.delete_object(obj)
        except Exception, e:
            module.fail_json(msg=e.message)
        else:
            results.append(result)

    num_deleted = results.count(True)

    EXIT_DICT['container'] = c.name
    EXIT_DICT['deleted'] = num_deleted
    EXIT_DICT['requested_deleted'] = objs

    if num_deleted:
        EXIT_DICT['changed'] = True

    if num_objs == num_deleted:
        EXIT_DICT['success'] = True
        EXIT_DICT['msg'] = "%s objects deleted" % num_deleted
    else:
        EXIT_DICT['msg'] = ("Error: only %s of %s objects "
                            "deleted" % (num_deleted, num_objs))
    module.exit_json(**EXIT_DICT)


def get_meta(module, cf, container, src, dest):
    """ Get metadata for a single file, comma-separated list, or entire
    container
    """
    c = _get_container(module, cf, container)

    objs = None
    if src and dest:
        module.fail_json(msg="Error: ambiguous instructions; files to be deleted "
                             "have been specified on both src and dest args")
    elif dest:
        objs = dest
    else:
        objs = src

    if objs:
        objs = objs.split(',')
        objs = map(str.strip, objs)
    else:
        objs = c.get_object_names()

    results = {}
    for obj in objs:
        try:
            meta = c.get_object(obj).get_metadata()
        except Exception, e:
            module.fail_json(msg=e.message)
        else:
            results[obj] = {k.split(META_PREFIX)[-1]: v for k, v in meta.iteritems()}

    EXIT_DICT['container'] = c.name
    if results:
        EXIT_DICT['meta_results'] = results
        EXIT_DICT['success'] = True
    module.exit_json(**EXIT_DICT)


def put_meta(module, cf, container, src, dest, meta, clear_meta):
    """ Set metadata on a container, single file, or comma-separated list.
    Passing a true value to clear_meta clears the metadata stored in Cloud
    Files before setting the new metadata to the value of "meta".
    """

    objs = None
    if src and dest:
        module.fail_json(msg="Error: ambiguous instructions; files to set meta"
                             " have been specified on both src and dest args")
    elif dest:
        objs = dest
    else:
        objs = src

    objs = objs.split(',')
    objs = map(str.strip, objs)

    c = _get_container(module, cf, container)

    results = []
    for obj in objs:
        try:
            result = c.get_object(obj).set_metadata(meta, clear=clear_meta)
        except Exception, e:
            module.fail_json(msg=e.message)
        else:
            results.append(result)

    EXIT_DICT['container'] = c.name
    EXIT_DICT['success'] = True
    if results:
        EXIT_DICT['changed'] = True
        EXIT_DICT['num_changed'] = True
    module.exit_json(**EXIT_DICT)


def delete_meta(module, cf, container, src, dest, meta):
    """ Removes metadata keys and values specified in meta, if any.  Deletes on
    all objects specified by src or dest (but not both), if any; otherwise it
    deletes keys on all objects in the container
    """
    objs = None
    if src and dest:
        module.fail_json(msg="Error: ambiguous instructions; meta keys to be "
                             "deleted have been specified on both src and dest"
                             " args")
    elif dest:
        objs = dest
    else:
        objs = src

    objs = objs.split(',')
    objs = map(str.strip, objs)

    c = _get_container(module, cf, container)

    results = []  # Num of metadata keys removed, not objects affected
    for obj in objs:
        if meta:
            for k, v in meta.items():
                try:
                    result = c.get_object(obj).remove_metadata_key(k)
                except Exception, e:
                    module.fail_json(msg=e.message)
                else:
                    results.append(result)
        else:
            try:
                o = c.get_object(obj)
            except pyrax.exc.NoSuchObject, e:
                module.fail_json(msg=e.message)

            for k, v in o.get_metadata().items():
                try:
                    result = o.remove_metadata_key(k)
                except Exception, e:
                    module.fail_json(msg=e.message)
                results.append(result)

    EXIT_DICT['container'] = c.name
    EXIT_DICT['success'] = True
    if results:
        EXIT_DICT['changed'] = True
        EXIT_DICT['num_deleted'] = len(results)
    module.exit_json(**EXIT_DICT)


def cloudfiles(module, container, src, dest, method, typ, meta, clear_meta,
               structure, expires):
        """ Dispatch from here to work with metadata or file objects """
        cf = pyrax.cloudfiles

        if typ == "file":
            if method == 'put':
                upload(module, cf, container, src, dest, meta, expires)

            elif method == 'get':
                download(module, cf, container, src, dest, structure)

            elif method == 'delete':
                delete(module, cf, container, src, dest)

        else:
            if method == 'get':
                get_meta(module, cf, container, src, dest)

            if method == 'put':
                put_meta(module, cf, container, src, dest, meta, clear_meta)

            if method == 'delete':
                delete_meta(module, cf, container, src, dest, meta)


def main():
    argument_spec = rax_argument_spec()
    argument_spec.update(
        dict(
            container=dict(required=True),
            src=dict(),
            dest=dict(),
            method=dict(default='get', choices=['put', 'get', 'delete']),
            type=dict(default='file', choices=['file', 'meta']),
            meta=dict(type='dict', default={}),
            clear_meta=dict(choices=BOOLEANS, default=False, type='bool'),
            structure=dict(choices=BOOLEANS, default=True, type='bool'),
            expires=dict(type='int'),
        )
    )

    module = AnsibleModule(
        argument_spec=argument_spec,
        required_together=rax_required_together()
    )

    container = module.params.get('container')
    src = module.params.get('src')
    dest = module.params.get('dest')
    method = module.params.get('method')
    typ = module.params.get('type')
    meta = module.params.get('meta')
    clear_meta = module.params.get('clear_meta')
    structure = module.params.get('structure')
    expires = module.params.get('expires')

    if clear_meta and not typ == 'meta':
        module.fail_json(msg='clear_meta can only be used when setting metadata')

    setup_rax_module(module, pyrax)
    cloudfiles(module, container, src, dest, method, typ, meta, clear_meta, structure, expires)


from ansible.module_utils.basic import *
from ansible.module_utils.rax import *

main()