Restore deleted files from AWS bucket

restore_aws_bucket.py

"""script to restore deleted objects of a AWS bucket
"""

import argparse
import json
import os
import subprocess as s
import tempfile


def get_bucket():
    parser = argparse.ArgumentParser()
    parser.add_argument('bucket', type=str, help="AWS Bucket that you want restored")
    return parser.parse_args().bucket


def restore_object(bucket, key, versions):
    backup = tempfile.mktemp()
    latest_version = sorted(versions, key=lambda x: x['LastModified'])[-1]['VersionId']
    s.check_output(f'aws s3api get-object --bucket {bucket} --key {key} --version-id {latest_version} {backup}', shell=True)
    s.check_output(f'aws s3api put-object --bucket {bucket} --key {key} --body {backup}', shell=True)
    os.remove(backup)
    print(f"restored {key}")


def restore_bucket(bucket):
    cmd = f"aws s3api list-object-versions --bucket {bucket}"
    process = s.Popen(cmd, shell=True, stdout=s.PIPE)
    out, err = process.communicate()
    data = json.loads(out.decode('utf-8'))

    # store all versions as a dict referenced by their key
    versions = {}
    for a in data["Versions"]:
        if a["Key"] not in versions:
            versions[a["Key"]] = []
        versions[a["Key"]].append(a)

    deleted_keys = sorted(set([object['Key'] for object in data['DeleteMarkers']]))

    for key in deleted_keys:
        p = s.Popen(f"aws s3api head-object --bucket {bucket} --key {key}", shell=True, stdout=s.PIPE, stderr=s.PIPE)
        out, err = p.communicate()

        if p.returncode == 0:
            print(f"key exists: {key}")
            continue

        restore_object(bucket, key, versions[key])


if __name__ == '__main__':
    restore_bucket(get_bucket())