s3-drop-versions

#!/usr/bin/env python3
#
# MIT License
#

import os, sys, time, signal, argparse, boto3, botocore, \
        unittest, concurrent.futures, shutil, uuid, time
from datetime import datetime, timezone
from dateutil.parser import parse

sys.stdout = open(sys.stdout.fileno(), mode='w', encoding='utf8', buffering=1)

args = None
executor = None
transfer = None
client = None
dest_bucket = ''
dest_prefix = ''
futures = {}


def signal_handler(signal, frame):
    print("Gracefully exiting...")
    if executor:
        print("Shutting down executor...")
        executor.shutdown(wait=False)
    if futures:
        for future in list(futures.keys()):
            print("Cancelling futures...")
            if not future.running():
                future.cancel()
                futures.pop(future, None)
    sys.exit(-1)

def s3_delte_key(obj):
    if not args.forreal or args.forreal != args.bucket:
        return
    r = client.delete_object(Bucket=args.bucket, Key=obj['Key'], VersionId=obj['VersionId'])
    return

del_buffer = []
deleted_size = 0
def mark_for_deletion(obj, cause_obj, cause_type):
    if cause_type == 'deleted':
        explanation = 'DELETED ON '+str(cause_obj["LastModified"])
    elif cause_type == 'overwritten':
        explanation = 'DEPRICATED ON '+str(cause_obj["LastModified"])
    elif cause_type == 'fully-wiped':
        explanation = "FULL WIPEOUT"
    elif cause_type == 'out-of-policy':
        explanation = 'RETAINED VERSION EXISTS FOR '+str(cause_obj["LastModified"])
    else:
        explanation = 'UNEXPECTED!!'

    if not args.forreal or args.forreal != args.bucket:
        action = "DRYRUN DELETING: "
    else:
        action = "DELETING: "
        future = executor.submit(s3_delte_key, obj)
        global futures
        futures[future] = obj
        global del_buffer
        #del_buffer.append({'Key':obj["Key"], 'VersionId':obj["VersionId"]})

    global deleted_size
    deleted_size += obj["Size"]

    print(action + obj["Key"]+ " " + str(obj["LastModified"])+" "+explanation)
    return

retained_size = 0
retained_sizes = {}
def retain(obj, explanation=""):
    global retained_size
    global retained_sizes

    if explanation not in retained_sizes:
        retained_sizes[explanation] = 0
    retained_sizes[explanation] += obj['Size']
    retained_size               += obj['Size']
    print("RETAIN:   "+obj["Key"]+" "+str(obj["LastModified"])+explanation)
    return


def do_clean_versions():
    global client
    client = boto3.client('s3')
    global transfer
    transfer = boto3.s3.transfer.S3Transfer(client)
    last_obj = {"Key":""}
    dmarker  = {"Key":""}
    global del_buffer
    del_buffer = []

    if args.debug: boto3.set_stream_logger('botocore')

    global executor
    executor = concurrent.futures.ThreadPoolExecutor(args.max_workers)

    objects = client.list_object_versions(Bucket=args.bucket, Prefix=args.prefix) # cannot pass VersionIdMarker on first call
    
    now = datetime.now(timezone.utc)    # TODO: not robust against faulty system time
                                        # but - we rely on AWS refusing our requests if theres is a drift > 15 minutes

    while (True):
        if not "Versions" in objects:
            print("Bucket not configured for versioning. Exiting.", file=sys.stderr)
            sys.exit(1)
        versions = objects['Versions']
        if 'DeleteMarkers' in objects:
            deletemarkers = objects['DeleteMarkers']
        else:
            deletemarkers = []

        while(versions):
            obj = versions.pop(0)

            if args.mode == 'delete-all':
                mark_for_deletion(obj, None, 'fully-wiped')
                continue

            elif args.mode == 'retain-latest':
                if obj["Key"] == last_obj["Key"]:
                    # This is an older version of the object and should be deleted
                    mark_for_deletion(obj, last_obj, 'overwritten')
                    continue

                # First time we see this key - check if it was deleted:
                while deletemarkers and (dmarker["Key"] < obj["Key"]):
                    dmarker = deletemarkers.pop(0)
                if dmarker["Key"] == obj["Key"] and dmarker["LastModified"] > obj["LastModified"]:
                    # The mostresent operation on this key was a delete - it should be deleted
                    mark_for_deletion(obj, dmarker, 'deleted')
                    last_obj = dmarker
                    continue

            elif args.mode == 'strategy':
                # Following the retention strategy:
                # * we don't look @ deletions
                # * we always retain the latest version (even if deleted)
                # * for versions 7 days or youger, we retain if the version is of a different day than the last retained one
                # * for versions 8-30 days old, we retain if the version is of a different week 
                # * for versions 30-365 days old, we retain if the version is of a different month
                # * for versions older than 365 days, we retain if the version is of a different year

                if obj["Key"] != last_obj["Key"]:
                    # First time we see it. Newest - retain even if deleted..
                    retain(obj, " (latest)")
                    last_obj = obj
                    continue

                age           = now - obj["LastModified"]
                retained_date = last_obj["LastModified"]
                version_date  = obj["LastModified"]

                if age.days < 7:
                    if (retained_date.day, retained_date.month, retained_date.year) != \
                        (version_date.day, version_date.month, version_date.year):
                        retain(obj, " (daily)")
                        last_obj = obj
                        continue
                elif age.days < 30:
                    if (retained_date.isocalendar()[1], retained_date.year) != \
                        (version_date.isocalendar()[1], version_date.year):
                        retain(obj, " (weekly)")
                        last_obj = obj
                        continue
                elif age.days < 365:
                    if (retained_date.month, retained_date.year) != \
                        (version_date.month, version_date.year):
                        retain(obj, " (month)")
                        last_obj = obj
                        continue
                else:
                    if retained_date.year != version_date.year:
                        retain(obj, " (yearly)")
                        last_obj = obj
                        continue

                # This is version does is extraneous acc. to retention strategy...
                mark_for_deletion(obj, last_obj, 'out-of-policy')
                continue

            # This object should be retained
            retain(obj)
            last_obj = obj

        if objects["IsTruncated"]:
            # More objects to be got
            objects = client.list_object_versions(Bucket=args.bucket, Prefix=args.prefix, \
                KeyMarker=objects["NextKeyMarker"], VersionIdMarker=objects["NextVersionIdMarker"])
        else:
            break

    #print(repr(del_buffer))

    for future in concurrent.futures.as_completed(futures):
        if future in futures:
            try:
                future.result()
                print("OK DELETE: "+repr(futures[future]))
            except Exception as ex:
                print('ERROR: %s --> %s"' % (ex, repr(futures[future])), file=sys.stderr)
            del(futures[future])


def bytes_to_gb(bytes):
    return(bytes/1024.0/1024.0/1024.0)

if __name__=='__main__':
    signal.signal(signal.SIGINT, signal_handler)

    parser = argparse.ArgumentParser()
    parser.add_argument('-b', '--bucket', help='s3 bucket to cleanup', required=True)
    parser.add_argument('-p', '--prefix', help='s3 path to restrict cleanup to', default="")
    parser.add_argument('-v', '--verbose', help='print verbose informations from s3 objects', action='store_true')
    parser.add_argument('--forreal', help="Dryrun is the default; Pass '--forreal <bucket>' to actually issue delete requests", default='')
    parser.add_argument('--debug', help='enable debug output', action='store_true')
    parser.add_argument('--max-workers', help='max number of concurrent delete requests', default=1000, type=int)
    parser.add_argument('--mode', help='What versions will be deleted', required=True, choices=['retain-all', 'retain-latest', 'delete-all', 'strategy'], default='retain-all')
    
    args = parser.parse_args()

    do_clean_versions()

    global deleted_size
    global retained_size
    global retained_sizes

    deleted_percent = float(deleted_size)/(retained_size+deleted_size)*100
    print("Retained Size: %.3f G" %(bytes_to_gb(retained_size)))
    print("Retained Size: %.3f G (%.3f%%)" %(bytes_to_gb(deleted_size), deleted_percent))

    for explanation in retained_sizes:
        print("Retained as %s: %.3f G"%(explanation, bytes_to_gb(retained_sizes[explanation])))

    sys.exit(0)