Source

cloudman / cm / util / misc.py

Full commit
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
#!/usr/bin/python
import logging, time, yaml
from boto.s3.key import Key
from boto.s3.acl import ACL
from boto.exception import S3ResponseError, EC2ResponseError
import subprocess
import threading
import datetime as dt
from tempfile import mkstemp
from shutil import move
from os import remove, close


log = logging.getLogger( 'cloudman' )

def load_yaml_string(str):
    ud = yaml.load(str)
    return ud

def load_yaml_file(filename):
    with open(filename) as ud_file:
        ud = yaml.load(ud_file)
    # log.debug("Loaded user data: %s" % ud)
    return ud
    
def dump_yaml_to_file(data, filename):
    with open(filename, 'w') as f:
        yaml.dump(data, f, default_flow_style=False)

def merge_yaml_objects(user, default):
    """Merge fields from user data (user) YAML object and default data (default)
    YAML object. If there are conflicts, value from the user data object are
    kept."""
    if isinstance(user,dict) and isinstance(default,dict):
        for k,v in default.iteritems():
            if k not in user:
                user[k] = v
            else:
                user[k] = merge_yaml_objects(user[k],v)
    return user

def shellVars2Dict(filename):
	'''Reads a file containing lines with <KEY>=<VALUE> pairs and turns it into a dict'''
	f = None
	try:
		f = open(filename, 'r')
	except IOError:
		return { }
	lines = f.readlines()
	result = { }

	for line in lines:
		parts = line.strip().partition('=')
		key = parts[0].strip()
		val = parts[2].strip()
		if key:
			result[key] = val
	return result

def formatDelta(delta):
    d = delta.days
    h = delta.seconds / 3600
    m = (delta.seconds % 3600) / 60
    s = delta.seconds % 60
    
    if d > 0:
        return '%sd %sh' % (d, h)
    elif h > 0:
        return '%sh %sm' % (h, m)
    else:
        return '%sm %ss' % (m, s)

def bucket_exists(s3_conn, bucket_name, validate=True):
    if bucket_name is not None:
        try:
            b = None
            b = s3_conn.lookup(bucket_name, validate=validate)
            if b is not None:
                # log.debug("Checking if bucket '%s' exists... it does." % bucket_name)
                return True
            else:
                log.debug("Checking if bucket '%s' exists... it does not." % bucket_name)
                return False
        except Exception, e:
            log.error("Failed to lookup bucket '%s': %s" % (bucket_name, e))
    else:
        log.error("Cannot lookup bucket with no name.")
        return False

def create_bucket(s3_conn, bucket_name):
    try:
        s3_conn.create_bucket(bucket_name)
        log.debug("Created bucket '%s'." % bucket_name)
    except S3ResponseError, e:
        log.error( "Failed to create bucket '%s': %s" % (bucket_name, e))
        return False
    return True

def get_bucket(s3_conn, bucket_name, validate=True):
    """Get handle to bucket"""
    b = None
    if bucket_exists(s3_conn, bucket_name, validate):
        for i in range(0, 5):
    		try:
    			b = s3_conn.get_bucket(bucket_name, validate=validate)
    			break
    		except S3ResponseError: 
    			log.error ( "Problem connecting to bucket '%s', attempt %s/5" % ( bucket_name, i+1 ) )
    			time.sleep(2)
    return b

def make_bucket_public(s3_conn, bucket_name, recursive=False):
    b = None
    b = get_bucket(s3_conn, bucket_name)
    if b is not None:
        try:
            b.make_public(recursive=recursive)
            log.debug("Bucket '%s' made public" % bucket_name)
            return True
        except S3ResponseError, e:
            log.error("Could not make bucket '%s' public: %s" % (bucket_name, e))
    return False

def make_key_public(s3_conn, bucket_name, key_name):
    b = None
    b = get_bucket(s3_conn, bucket_name)
    if b is not None:
        try:
            k = Key(b, key_name)
            if k.exists():
                k.make_public()
                log.debug("Key '%s' made public" % key_name)
                return True
        except S3ResponseError, e:
            log.error("Could not make key '%s' public: %s" % (key_name, e))
    return False

def add_bucket_user_grant(s3_conn, bucket_name, permission, cannonical_ids, recursive=False):
    """
    Boto wrapper that provides a quick way to add a canonical
    user grant to a bucket. 
    
    :type permission: string
    :param permission: The permission being granted. Should be one of:
                       (READ, WRITE, READ_ACP, WRITE_ACP, FULL_CONTROL).
    
    :type user_id: list of strings
    :param cannonical_ids: A list of strings with canonical user ids associated 
                        with the AWS account your are granting the permission to.
    
    :type recursive: boolean
    :param recursive: A boolean value to controls whether the command
                      will apply the grant to all keys within the bucket
                      or not.
    """
    b = None
    b = get_bucket(s3_conn, bucket_name)
    if b is not None:
        try:
            for c_id in cannonical_ids:
                log.debug("Adding '%s' permission for bucket '%s' for users '%s'" % (permission, bucket_name, c_id))
                b.add_user_grant(permission, c_id, recursive)
            return True
        except S3ResponseError, e:
            log.error("Could not add permission '%s' for bucket '%s': %s" % (permission, bucket_name, e))
    return False

def add_key_user_grant(s3_conn, bucket_name, key_name, permission, cannonical_ids):
    """
    Boto wrapper that provides a quick way to add a canonical
    user grant to a key. 
    
    :type permission: string
    :param permission: Name of the bucket where the key resides
    
    :type permission: string
    :param permission: Name of the key to add the permission to
    
    :type permission: string
    :param permission: The permission being granted. Should be one of:
                       (READ, WRITE, READ_ACP, WRITE_ACP, FULL_CONTROL).
    
    :type user_id: list of strings
    :param cannonical_ids: A list of strings with canonical user ids associated 
                        with the AWS account your are granting the permission to.
    """
    b = None
    b = get_bucket(s3_conn, bucket_name)
    if b is not None:
        try:
            k = Key(b, key_name)
            if k.exists():
                for c_id in cannonical_ids:
                    log.debug("Adding '%s' permission for key '%s' for user '%s'" % (permission, key_name, c_id))
                    k.add_user_grant(permission, c_id)
                return True
        except S3ResponseError, e:
            log.error("Could not add permission '%s' for bucket '%s': %s" % (permission, bucket_name, e))
    return False

def get_list_of_bucket_folder_users(s3_conn, bucket_name, folder_name, exclude_power_users=True):
    """
    Retrieve a list of users that are associated with a key in a folder (i.e., prefix) 
    in the provided bucket and have READ grant. Note that this method assumes all 
    of the keys in the given folder have the same ACL and the method thus looks 
    only at the very first key in the folder. Also, any users
    
    :type s3_conn: boto.s3.connection.S3Connection
    :param s3_conn: Established boto connection to S3 that has access to bucket_name
    
    :type bucket_name: string
    :param bucket_name: Name of the bucket in which the given folder/prefix is 
                        stored
    
    :type folder_name: string
    :param folder_name: Allows limit the listing of keys in the given bucket to a 
                   particular prefix. This has the effect of iterating through 
                   'folders' within a bucket. For example, if you call the method 
                   with folder_name='/foo/' then the iterator will only cycle
                   through the keys that begin with the string '/foo/'.
                   A valid example would be 'shared/2011-03-31--19-43/'
    
    :type exclude_power_users: boolean
    :param exclude_power_users: If True, folder users with FULL_CONTROL grant 
                   are not included in the folder user list
    """
    users=[] # Current list of users retrieved from folder's ACL
    key_list = None
    key_acl = None
    b = None
    b = get_bucket(s3_conn, bucket_name)
    if b is not None:
        try:
            key_list = b.get_all_keys(prefix=folder_name, delimiter='/')
            # for k in key_list:
            #     print k.name#, k.get_acl().acl.grants[0].type
            if len(key_list) > 0:
                key = key_list[0] # Just get one key assuming all keys will have the same ACL
                key_acl = key.get_acl()
            if key_acl:
                power_users = []
                for grant in key_acl.acl.grants:
                    # log.debug("folder_name: %s, %s, %s, %s, %s." % (folder_name, key.name, grant.type, grant.display_name, grant.permission))
                    if grant.permission == 'FULL_CONTROL':
                        power_users.append(grant.display_name)
                    if grant.type == 'Group' and 'Group' not in users:
                        # Group grants (i.e., public) are simply listed as Group under grant.type so catch that
                        users.append(u'Group')
                    elif grant.permission == 'READ' and grant.type != 'Group' and grant.display_name not in users:
                        users.append(grant.display_name)
                # Users w/ FULL_CONTROL are optionally not included in the folder user list
                if exclude_power_users:
                    for pu in power_users:
                        if pu in users:
                            users.remove(pu)
        except S3ResponseError, e:
            log.error("Error getting list of folder '%s' users for bucket '%s': %s" % (folder_name, bucket_name, e))
    # log.debug("List of users for folder '%s' in bucket '%s': %s" % (folder_name, bucket_name, users))
    return users

def get_users_with_grant_on_only_this_folder(s3_conn, bucket_name, folder_name):
    """
    This method is used when dealing with bucket permissions of shared instances. 
    The method's intent is to isolate the set of users that have (READ) grant on 
    a given folder and no other (shared) folder within the given bucket.
    Obtained results can then be used to set the permissions on the bucket root.
    
    See also: get_list_of_bucket_folder_users
    
    :type s3_conn: boto.s3.connection.S3Connection
    :param s3_conn: Established boto connection to S3 that has access to bucket_name
    
    :type bucket_name: string
    :param bucket_name: Name of the bucket in which the given folder is stored
    
    :type folder_name: string
    :param folder_name: Name of the (shared) folder whose grants will be examined
                        and compared to other (shared) folders in the same bucket.
                        A valid example would be 'shared/2011-03-31--19-43/'
    """
    users_with_grant = [] # List of users with grant on given folder and no other (shared) folder in bucket
    other_users = [] # List of users on other (shared) folders in given bucket
    folder_users = get_list_of_bucket_folder_users(s3_conn, bucket_name, folder_name)
    # log.debug("List of users on to-be-deleted shared folder '%s': %s" % (folder_name, folder_users))
    b = None
    b = get_bucket(s3_conn, bucket_name)
    if b is not None:
        try:
            # Get list of shared folders in given bucket
            folder_list = b.get_all_keys(prefix='shared/', delimiter='/')
            # Inspect each shared folder's user grants and create a list of all users
            # with grants on those folders
            for f in folder_list:
                if f.name != folder_name:
                    fu = get_list_of_bucket_folder_users(s3_conn, bucket_name, f.name)
                    for u in fu:
                        if u not in other_users:
                            other_users.append(u)
            # log.debug("List of users on other shared folders: %s" % other_users)
            # Find list of users in that have grants only on 'folder_name' and no
            # other shared folder in the given bucket
            for u in folder_users:
                if u not in other_users:
                    users_with_grant.append(u)
        except S3ResponseError, e:
            log.error("Error isolating list of folder '%s' users for bucket '%s': %s" % (folder_name, bucket_name, e))
    log.debug("List of users whose bucket grant is to be removed because shared folder '%s' is being deleted: %s" \
        % (folder_name, users_with_grant))
    return users_with_grant

def adjust_bucket_ACL(s3_conn, bucket_name, users_whose_grant_to_remove):
    """
    Adjust the ACL on given bucket and remove grants for all the mentioned users.
    
    :type s3_conn: boto.s3.connection.S3Connection
    :param s3_conn: Established boto connection to S3 that has access to bucket_name
    
    :type bucket_name: string
    :param bucket_name: Name of the bucket for which to adjust the ACL
    
    :type users_whose_grant_to_remove: list
    :param users_whose_grant_to_remove: List of user names (as defined in bucket's
                initial ACL, e.g., ['afgane', 'cloud']) whose grant is to be revoked.
    """
    bucket = get_bucket(s3_conn, bucket_name)
    if bucket:
        try:
            grants_to_keep = []
            # log.debug("All grants on bucket '%s' are following" % bucket_name)
            # Compose list of grants on the bucket that are to be kept, i.e., siphon 
            # through the list of grants for bucket's users and the list of users
            # whose grant to remove and create a list of bucket grants to keep
            for g in bucket.get_acl().acl.grants:
                # log.debug("Grant -> permission: %s, user name: %s, grant type: %s" % (g.permission, g.display_name, g.type))
                # Public (i.e., group) permissions are kept under 'type' field so check that first
                if g.type == 'Group' and 'Group' in users_whose_grant_to_remove:
                    pass
                elif g.display_name not in users_whose_grant_to_remove:
                    grants_to_keep.append(g)
            # Manipulate bucket's ACL now
            bucket_policy = bucket.get_acl() # Object for bucket's current policy (which holds the ACL)
            acl = ACL() # Object for bucket's to-be ACL
            # Add all of the exiting (i.e., grants_to_keep) grants to the new ACL object
            for gtk in grants_to_keep:
                acl.add_grant(gtk)
            # Update the policy and set bucket's ACL
            bucket_policy.acl = acl
            bucket.set_acl(bucket_policy)
            # log.debug("List of kept grants for bucket '%s'" % bucket_name)
            # for g in bucket_policy.acl.grants:
            #     log.debug("Grant -> permission: %s, user name: %s, grant type: %s" % (g.permission, g.display_name, g.type))
            log.debug("Removed grants on bucket '%s' for these users: %s" % (bucket_name, users_whose_grant_to_remove))
            return True
        except S3ResponseError, e:
            log.error("Error adjusting ACL for bucket '%s': %s" % (bucket_name, e))
    return False

def file_exists_in_bucket(s3_conn, bucket_name, remote_filename):
    """Check if remote_filename exists in bucket bucket_name.
    :rtype: bool
    :return: True if remote_filename exists in bucket_name
             False otherwise
    """
    b = None
    b = get_bucket(s3_conn, bucket_name)
    if b is not None:
        try:
    		k = Key(b, remote_filename)
    		if k.exists():
    		    return True
    	except S3ResponseError:
    	    log.debug("Key '%s' in bucket '%s' does not exist." % (remote_filename, bucket_name))
    return False

def get_file_from_bucket( conn, bucket_name, remote_filename, local_file, validate=True ):
    if bucket_exists(conn, bucket_name, validate):
        b = get_bucket(conn, bucket_name, validate)
        k = Key(b, remote_filename)
        try:
            k.get_contents_to_filename(local_file)
            log.info( "Retrieved file '%s' from bucket '%s' to '%s'." \
                         % (remote_filename, bucket_name, local_file))
        except S3ResponseError, e:
            log.debug( "Failed to get file '%s' from bucket '%s': %s" % (remote_filename, bucket_name, e))
            return False
    else:
      log.debug("Bucket '%s' does not exist, did not get remote file '%s'" % (bucket_name, remote_filename))
      return False
    return True

def save_file_to_bucket( conn, bucket_name, remote_filename, local_file ):
    b = None
    b = get_bucket(conn, bucket_name)
    if b is not None:
		k = Key( b, remote_filename )
		try:
		    k.set_contents_from_filename( local_file )
		    log.info( "Saved file '%s' to bucket '%s'" % ( remote_filename, bucket_name ) )
		    # Store some metadata (key-value pairs) about the contents of the file being uploaded
		    k.set_metadata('date_uploaded', dt.datetime.utcnow())
		except S3ResponseError, e:
		     log.error( "Failed to save file local file '%s' to bucket '%s' as file '%s': %s" % ( local_file, bucket_name, remote_filename, e ) )
		     return False
		return True
    else:
        log.debug("Could not connect to bucket '%s'; remote file '%s' not saved to the bucket" % (bucket_name, remote_filename))
        return False

def copy_file_in_bucket(s3_conn, src_bucket_name, dest_bucket_name, orig_filename, copy_filename, preserve_acl=True, validate=True):
    b = get_bucket(s3_conn, src_bucket_name, validate)
    if b:
        try:
            log.debug("Establishing handle with key object '%s'" % orig_filename)
            k = Key(b, orig_filename)
            log.debug("Copying file '%s/%s' to file '%s/%s'" % (src_bucket_name, orig_filename, dest_bucket_name, copy_filename))
            k.copy(dest_bucket_name, copy_filename, preserve_acl=preserve_acl)
            return True
        except S3ResponseError, e: 
            log.debug("Error copying file '%s/%s' to file '%s/%s': %s" % (src_bucket_name, orig_filename, dest_bucket_name, copy_filename, e))
    return False

def delete_file_from_bucket(conn, bucket_name, remote_filename):
    b = None
    b = get_bucket(conn, bucket_name)
    if b is not None:
        try:
            k = Key( b, remote_filename )
            log.debug( "Deleting key object '%s' from bucket '%s'" % (remote_filename, bucket_name))
            k.delete()
            return True
    	except S3ResponseError, e:
    	    log.error("Error deleting key '%s' from bucket '%s': %s" % (remote_filename, bucket_name, e))
    return False

def delete_bucket(conn, bucket_name):
    """Delete given bucket. This method will iterate through all the keys in
    the given bucket first and delete them. Finally, the bucket will be deleted."""
    try:
        b = None
        b = get_bucket(conn, bucket_name)
        if b is not None:
            keys = b.get_all_keys()
            for key in keys:
                key.delete()
            b.delete()
            log.info("Successfully deleted cluster bucket '%s'" % bucket_name)
    except S3ResponseError, e:
        log.error("Error deleting bucket '%s': %s" % (bucket_name, e))
    return True

def get_file_metadata(conn, bucket_name, remote_filename, metadata_key):
    """ Get metadata value for the given key. If bucket or remote_filename is not 
    found, the method returns None.    
    """
    log.debug("Getting metadata '%s' for file '%s' from bucket '%s'" % (metadata_key, remote_filename, bucket_name))
    b = None
    b = get_bucket(conn, bucket_name)
    if b is not None:
        k = b.get_key(remote_filename)
        if k and metadata_key:
            return k.get_metadata(metadata_key)
    return None 

def set_file_metadata(conn, bucket_name, remote_filename, metadata_key, metadata_value):
    """ Set metadata key-value pair for the remote file. 
    :rtype: bool
    :return: If specified bucket and remote_filename exist, return True.
             Else, return False
    """
    log.debug("Setting metadata '%s' for file '%s' in bucket '%s'" % (metadata_key, remote_filename, bucket_name))

    b = None
    b = get_bucket(conn, bucket_name)
    if b is not None:
        k = b.get_key(remote_filename)
        if k and metadata_key:
            # Simply setting the metadata through set_metadata does not work.
            # Instead, must create in-place copy of the file with altered metadata:
            # http://groups.google.com/group/boto-users/browse_thread/thread/9968d3fc4fc18842/29c680aad6e31b3e#29c680aad6e31b3e 
            try:
                k.copy(bucket_name, remote_filename, metadata={metadata_key:metadata_value}, preserve_acl=True)
                return True
            except Exception, e:
                log.debug("Could not set metadata for file '%s' in bucket '%s': %e" % (remote_filename, bucket_name, e))
    return False

def check_process_running(proc_in):
    ps = subprocess.Popen("ps ax", shell=True, stdout=subprocess.PIPE)
    lines = str(ps.communicate()[0], 'utf-8')
    for line in lines.split('\n'):
        if line.find(proc_in) != -1:
            log.debug(line)
    ps.stdout.close()
    ps.wait()
    
def get_volume_size(ec2_conn, vol_id):
    log.debug("Getting size of volume '%s'" % vol_id) 
    try: 
        vol = ec2_conn.get_all_volumes(vol_id)
    except EC2ResponseError, e:
        log.error("Volume ID '%s' returned an error: %s" % (vol_id, e))
        return 0 

    if len(vol) > 0: 
        # We're only looking for the first vol bc. that's all we asked for
        return vol[0].size
    else:
        return 0

def run(cmd, err='cmd failed', ok='cmd OK'):
    """ Convenience method for executing a shell command. """
    process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = process.communicate()
    if process.returncode == 0:
        log.debug(ok)
        return True
    else:
        log.error("%s, running command '%s' returned code '%s' and following stderr: '%s'" % (err, cmd, process.returncode, stderr))
        return False

def replace_string(file_name, pattern, subst):
    """Replace string in a file
    :type file_name: str
    :param file_name: Full path to the file where the string is to be replaced

    :type pattern: str
    :param pattern: String pattern to search for
    
    :type subst: str
    :param subst: String pattern to replace search pattern with 
    """
    # Create temp file
    fh, abs_path = mkstemp()
    new_file = open(abs_path,'w')
    old_file = open(file_name)
    for line in old_file:
        new_file.write(line.replace(pattern, subst))
    # Close temp file
    new_file.close()
    close(fh)
    old_file.close()
    # Remove original file
    remove(file_name)
    # Move new file
    move(abs_path, file_name)


class Sleeper( object ):
    """
    Provides a 'sleep' method that sleeps for a number of seconds *unless*
    the notify method is called (from a different thread).
    """
    def __init__( self ):
        self.condition = threading.Condition()
    def sleep( self, seconds ):
        self.condition.acquire()
        self.condition.wait( seconds )
        self.condition.release()
    def wake( self ):
        self.condition.acquire()
        self.condition.notify()
        self.condition.release()