ansible/library/cloud/gc_storage

#!/usr/bin/python
# 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: gc_storage
short_description: This module help users manage objects/buckets in GCS (Google Cloud Storage). 
description:
    - This module allows users to manage their objects/buckets in GCS in an idempotent way. Users can do operations like PUT, GET and set some pre-canned permissions on those objects.Users can also retrieve the urls to their objects to be used in playbooks. They can also use the module to get the string contents of the objects. Please note that this module is compatible with API version 1 of GCS. so it expects the user to set the default project before using this module, Have a lookt at "https://developers.google.com/storage/docs/reference/v1/apiversion1" to set the default project.

options:
  bucket:
    description:
      - Bucket name. 
    required: true
    default: null 
    aliases: []
  object:
    description:
      - Keyname of the object inside the bucket. Can be also be used to create "virtual directories", see examples.
    required: false
    default: null
    aliases: []
  src:
    description:
      - The source file path when performing a PUT operation.
    required: false
    default: null
    aliases: []
  dest:
    description:
      - The destination file path when downloading an object/key with a GET operation.
    required: false
    aliases: []
  overwrite:
    description:
      - Force overwrite either locally on the filesystem or remotely with the object/key. Used with PUT and GET operations.
    required: false
    default: false
  permission:
    description:
      - This option let's the user set the canned permissions on the object/bucket that are created. The permissions that can be set are 'private', 'public-read', 'authenticated-read'.
    required: false
    default: private 
  expiry:
    description:
      - Time limit (in seconds) for the URL generated and returned by GCA when performing a mode=put or mode=geturl operation. This url is only avaialbe when public-read is the acl for the object.
    required: false
    default: null
    aliases: []
  mode:
    description:
      - Switches the module behaviour between put (upload), get (download), geturl (return download url) , getstr (download object as string), create (bucket) and delete (bucket). 
    required: true
    default: null
    aliases: []
  gcs_secret_key:
    description:
      - GCS secret key. If not set then the value of the GCS_SECRET_KEY environment variable is used. 
    required: true
    default: null
  gcs_access_key:
    description:
      - GCS access key. If not set then the value of the GCS_ACCESS_KEY environment variable is used.
    required: true
    default: null

requirements: [ "boto 2.9+" ]

author: benno@ansibleworks.com Note. Most of the code has been taken from the S3 module.

'''

EXAMPLES = '''
# Simple PUT operation
- gcs: bucket=mybucket object=key.txt src=/usr/local/myfile.txt mode=put permission=public-read
# Simple GET operation
- gs: bucket=mybucket object=key.txt dest=/usr/local/myfile.txt mode=get
# GET/download and overwrite local file (trust remote)
- gs: bucket=mybucket object=key.txt dest=/usr/local/myfile.txt mode=get overwrite=true
# PUT/upload and overwrite remote file (trust local)
- gs: bucket=mybucket object=key.txt src=/usr/local/myfile.txt mode=put overwrite=true
# Download an object as a string to use else where in your playbook
- gs: bucket=mybucket object=key.txt mode=getstr
# Create an empty bucket
- gs: bucket=mybucket mode=create
# Create a bucket with key as directory
- gs: bucket=mybucket object=/my/directory/path mode=create
# Delete a bucket and all contents
- gs: bucket=mybucket mode=delete
'''

import sys
import os
import urlparse
import hashlib

try:
    import boto
except ImportError:
    print "failed=True msg='boto 2.9+ required for this module'"
    sys.exit(1)

def grant_check(module, gs, obj):
    try:
        acp = obj.get_acl()
        if module.params.get('permission') == 'public-read':
            grant = [ x for x in acp.entries.entry_list if x.scope.type == 'AllUsers']
            if not grant:
                obj.set_acl('public-read')
                module.exit_json(changed=True, result="The object's permission as been set to public-read") 
        if module.params.get('permission') == 'authenticated-read':
            grant = [ x for x in acp.entries.entry_list if x.scope.type == 'AllAuthenticatedUsers']
            if not grant:
                obj.set_acl('authenticated-read')
                module.exit_json(changed=True, result="The object's permission as been set to authenticated-read") 
    except gs.provider.storage_response_error, e:
        module.fail_json(msg= str(e))
    return True

        

def key_check(module, gs, bucket, obj):
    try:
        bucket = gs.lookup(bucket)
        key_check = bucket.get_key(obj)
    except gs.provider.storage_response_error, e:
        module.fail_json(msg= str(e))
    if key_check:
        grant_check(module, gs, key_check)
        return True
    else:
        return False

def keysum(module, gs, bucket, obj):
    bucket = gs.lookup(bucket)
    key_check = bucket.get_key(obj)
    if key_check:
        md5_remote = key_check.etag[1:-1]
        etag_multipart = md5_remote.find('-')!=-1 #Check for multipart, etag is not md5
        if etag_multipart is True:
            module.fail_json(msg="Files uploaded with multipart of gs are not supported with checksum, unable to compute checksum.")
    return md5_remote

def bucket_check(module, gs, bucket):
    try:
        result = gs.lookup(bucket)
    except gs.provider.storage_response_error, e:
        module.fail_json(msg= str(e))
    if result:
        grant_check(module, gs, result)
        return True
    else:
        return False

def create_bucket(module, gs, bucket):
    try:
        bucket = gs.create_bucket(bucket)
        bucket.set_acl(module.params.get('permission'))
    except gs.provider.storage_response_error, e:
        module.fail_json(msg= str(e))
    if bucket:
        return True

def delete_bucket(module, gs, bucket):
    try:
        bucket = gs.lookup(bucket)
        bucket_contents = bucket.list()
        for key in bucket_contents:
            bucket.delete_key(key.name)
        bucket.delete()
        return True
    except gs.provider.storage_response_error, e:
        module.fail_json(msg= str(e))

def delete_key(module, gs, bucket, obj):
    try:
        bucket = gs.lookup(bucket)
        bucket.delete_key(obj)
        module.exit_json(msg="Object deleted from bucket ", changed=True)
    except gs.provider.storage_response_error, e:
        module.fail_json(msg= str(e))
 
def create_dirkey(module, gs, bucket, obj):
    try:
        bucket = gs.lookup(bucket)
        key = bucket.new_key(obj)
        key.set_contents_from_string('')
        module.exit_json(msg="Virtual directory %s created in bucket %s" % (obj, bucket.name), changed=True)
    except gs.provider.storage_response_error, e:
        module.fail_json(msg= str(e))

def upload_file_check(src):
    if os.path.exists(src):
        file_exists is True
    else:
        file_exists is False
    if os.path.isdir(src):
        module.fail_json(msg="Specifying a directory is not a valid source for upload.", failed=True)
    return file_exists

def path_check(path):
    if os.path.exists(path):
        return True 
    else:
        return False

def upload_gsfile(module, gs, bucket, obj, src, expiry):
    try:
        bucket = gs.lookup(bucket)
        key = bucket.new_key(obj)  
        key.set_contents_from_filename(src)
        key.set_acl(module.params.get('permission'))
        url = key.generate_url(expiry)
        module.exit_json(msg="PUT operation complete", url=url, changed=True)
    except gs.provider.storage_copy_error, e:
        module.fail_json(msg= str(e))

def download_gsfile(module, gs, bucket, obj, dest):
    try:
        bucket = gs.lookup(bucket)
        key = bucket.lookup(obj)
        key.get_contents_to_filename(dest)
        module.exit_json(msg="GET operation complete", changed=True)
    except gs.provider.storage_copy_error, e:
        module.fail_json(msg= str(e))

def download_gsstr(module, gs, bucket, obj):
    try:
        bucket = gs.lookup(bucket)
        key = bucket.lookup(obj)
        contents = key.get_contents_as_string()
        module.exit_json(msg="GET operation complete", contents=contents, changed=True)
    except gs.provider.storage_copy_error, e:
        module.fail_json(msg= str(e))

def get_download_url(module, gs, bucket, obj, expiry):
    try:
        bucket = gs.lookup(bucket)
        key = bucket.lookup(obj)
        url = key.generate_url(expiry)
        module.exit_json(msg="Download url:", url=url, expiry=expiry, changed=True)
    except gs.provider.storage_response_error, e:
        module.fail_json(msg= str(e))

def handle_get(module, gs, bucket, obj, overwrite, dest):
    md5_remote = keysum(module, gs, bucket, obj)
    md5_local = hashlib.md5(open(dest, 'rb').read()).hexdigest()
    if md5_local == md5_remote:
        module.exit_json(msg="Local and remote object are identical", changed=False)
    else:
        if overwrite is True:
            download_gsfile(module, gs, bucket, obj, dest)
        else:
            module.fail_json(msg="Destination file exists and is diffrent. Use overwrite parameter to force download.", failed=True)

def handle_put(module, gs, bucket, obj, overwrite, src, expiry):
    # Lets check to see if bucket exists to get ground truth.
    bucket_rc = bucket_check(module, gs, bucket)
    key_rc    = key_check(module, gs, bucket, obj)

    # Lets check key state. Does it exist and if it does, compute the etag md5sum.
    if bucket_rc and key_rc:
        md5_remote = keysum(module, gs, bucket, obj)
        md5_local = hashlib.md5(open(src, 'rb').read()).hexdigest()
        if md5_local == md5_remote:
            module.exit_json(msg="Local and remote object are identical.", changed=False)
        else:
            if overwrite:
                upload_gsfile(module, gs, bucket, obj, src, expiry)
            else:
                module.exit_json(msg="WARNING: Checksums do not match. Use overwrite parameter to force upload.", failed=True)
                                                                                                            
    if not bucket_rc:      
        create_bucket(module, gs, bucket)
        upload_gsfile(module, gs, bucket, obj, src, expiry)

    # If bucket exists but key doesn't, just upload.
    if bucket_rc and not key_rc:
            upload_gsfile(module, gs, bucket, obj, src, expiry)
    
def handle_delete(module, gs, bucket, obj):
    if bucket and not obj:
        if bucket_check(module, gs, bucket):
            module.exit_json(msg="Bucket %s and all keys have been deleted."%bucket, changed=delete_bucket(module, gs, bucket))
        else:
            module.exit_json(msg="Bucket does not exist.", changed=False)
    if bucket and obj:
        if bucket_check(module, gs, bucket):
            if key_check(module, gs, bucket, obj):
                module.exit_json(msg="Object has been deleted.", changed=delete_key(module, gs, bucket, obj))
            else:
                module.exit_json(msg="Object does not exists.", changed=False)
        else:
            module.exit_json(msg="Bucket does not exist.", changed=False)
    else:
        module.fail_json(msg="Bucket or Bucket & object  parameter is required.", failed=True)
 
def handle_create(module, gs, bucket, obj):
    if bucket and not obj: 
        if bucket_check(module, gs, bucket):
            module.exit_json(msg="Bucket already exists.", changed=False)
        else:
            module.exit_json(msg="Bucket created succesfully", changed=create_bucket(module, gs, bucket))
    if bucket and obj:
        if bucket_check(module, gs, bucket):
            if obj.endswith('/'):
                dirobj = obj
            else:
                dirobj = obj + "/"
            if key_check(module, gs, bucket, dirobj):
                module.exit_json(msg="Bucket %s and key %s already exists."% (bucket, obj), changed=False)
            else:      
                create_dirkey(module, gs, bucket, dirobj)
        else:
            create_bucket(module, gs, bucket)
            create_dirkey(module, gs, bucket, dirobj)

def main():
    module = AnsibleModule(
        argument_spec = dict(
            bucket         = dict(required=True),
            object         = dict(default=None),
            src            = dict(default=None),
            dest           = dict(default=None),
            expiry         = dict(default=600, aliases=['expiration']),
            mode           = dict(choices=['get', 'put', 'delete', 'create', 'geturl', 'getstr'], required=True),
            permission     = dict(choices=['private', 'public-read', 'authenticated-read'], default='private'),
            gs_secret_key  = dict(no_log=True, required=True),
            gs_access_key  = dict(required=True),
            overwrite      = dict(default=False, type='bool'),
        ),
    )

    bucket        = module.params.get('bucket')
    obj           = module.params.get('object')
    src           = module.params.get('src')
    dest          = module.params.get('dest')
    if dest:
        dest      = os.path.expanduser(dest)
    mode          = module.params.get('mode')
    expiry        = module.params.get('expiry')
    gs_secret_key = module.params.get('gs_secret_key')
    gs_access_key = module.params.get('gs_access_key')
    overwrite     = module.params.get('overwrite')

    if mode == 'put':
        if not src or not object:
            module.fail_json(msg="When using PUT, src, bucket, object are mandatory paramters") 
    if mode == 'get':
        if not dest or not object:
            module.fail_json(msg="When using GET, dest, bucket, object are mandatory paramters") 
    if obj:
        obj = os.path.expanduser(module.params['object'])

    # ---------------- Get the Conection object for GCS operations --------------------
    try:
        gs = boto.connect_gs(gs_access_key, gs_secret_key)
    except boto.exception.NoAuthHandlerFound, e:
        module.fail_json(msg = str(e))
 
    #-------------- Code for GET Operation -------------------------
    if mode == 'get':
        if not bucket_check(module, gs, bucket) or not key_check(module, gs, bucket, obj):
            module.fail_json(msg="Target bucket/key cannot be found", failed=True)
        if not path_check(dest):
            download_gsfile(module, gs, bucket, obj, dest)
        else:
            handle_get(module, gs, bucket, obj, overwrite, dest)
        
    # ----------------------Code for PUT operation ---------------        
    if mode == 'put':
        if not path_check(src):
            module.fail_json(msg="Local object for PUT does not exist", failed=True)
        handle_put(module, gs, bucket, obj, overwrite, src, expiry)

    # Support for deleting an object if we have both params.  
    if mode == 'delete':
        handle_delete(module, gs, bucket, obj)
    
    # ---------- Code for mode is create -------------------------
    if mode == 'create':
        handle_create(module, gs, bucket, obj)
    
    # ------------- Get the url of an object --------------------------
    if mode == 'geturl':
        if bucket and obj:
            if bucket_check(module, gs, bucket) and key_check(module, gs, bucket, obj):
                get_download_url(module, gs, bucket, obj, expiry)
            else:
                module.fail_json(msg="Key/Bucket doesnt exist", failed=True)
        else:
            module.fail_json(msg="Bucket and Object parameters must be set", failed=True)

    # --------------------------- Get the String contents of an Object -------------------------
    if mode == 'getstr':
        if bucket and obj:
            if bucket_check(module, gs, bucket) and key_check(module, gs, bucket, obj):
                download_gsstr(module, gs, bucket, obj)
            else:
                module.fail_json(msg="Key/Bucket doesnt exists", failed=True)
        else:
            module.fail_json(msg="Bucket and Object parameters must be set", failed=True)


# this is magic, see lib/ansible/module_common.py
#<<INCLUDE_ANSIBLE_MODULE_COMMON>>

main()