<commit_before>#!/usr/bin/python3
"""
Get current ACL of all objects in given S3 bucket or set them to private or revert back.
Script supports 3 operations
1- getacl
2- setaclprivate
3- revertacl

It saves current ACL in a file named bucketname.txt for updating or reverting purposes.

python s3_acl.py --bucketname <name-of-bucket> --operation getacl

Should assume role to run this script.
"""


import boto3
from botocore.exceptions import ClientError
import backoff
import sys
import json
import click
import logging

MAX_TRIES = 5
region = "us-east-1"
# Set logging configuration
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# create file handler that logs messages
filehandler = logging.FileHandler('s3_acl.log')
filehandler.setLevel(logging.INFO)
formatter = logging.Formatter(
    '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
filehandler.setFormatter(formatter)
# add the handlers to logger
logger.addHandler(filehandler)


class S3BotoWrapper:
    def __init__(self, **kwargs):
        self.client = boto3.client("s3", **kwargs)

    @backoff.on_exception(backoff.expo,
                          ClientError,
                          max_tries=MAX_TRIES)
    def get_object(self, *args, **kwargs):
        return self.client.list_objects_v2(*args, **kwargs)

    @backoff.on_exception(backoff.expo,
                          ClientError,
                          max_tries=MAX_TRIES)
    def get_acl(self, *args, **kwargs):
        return self.client.get_object_acl(*args, **kwargs)

    @backoff.on_exception(backoff.expo,
                          ClientError,
                          max_tries=MAX_TRIES)
    def put_acl(self, *args, **kwargs):
        return self.client.put_object_acl(*args, **kwargs)


def get_all_s3_keys(s3_bucket, region):
    """Get a list of all keys in an S3 bucket."""
    keys = []

    kwargs = {'Bucket': s3_bucket}
    while True:
        s3_client = S3BotoWrapper(region_name=region)
        resp = s3_client.get_object(**kwargs)
        for obj in resp['Contents']:
            if obj['Key'][-1] == "/": # Filter out directories, you can add more filters here if required.
                continue
            else:
                keys.append(obj['Key'])
        try:
            kwargs['ContinuationToken'] = resp['NextContinuationToken']
        except KeyError:
            break
    return keys


def set_acl_private(acl_list, bucket_name):
    s3_client = S3BotoWrapper(region_name=region)
    for item in acl_list:
        for key, value in item.items():
            try:
                s3_client.put_acl(
                    ACL='private',
                    Bucket=bucket_name,
                    Key=key,
                )
            except ClientError as e:
                if e.response['Error']['Code'] == 'NoSuchKey':
                    logger.warning("No such key in S3: " + key)  # Will send the errors to the file
                else:
                    logger.error(("Unexpected error :{}".format(e)))
                    sys.exit(1)


def revert_s3_acl(acl_list, bucket_name):
    s3_client = S3BotoWrapper(region_name=region)
    for item in acl_list:
        for key, value in item.items():
            try:
                value.pop('ResponseMetadata', None)
                s3_client.put_acl(
                    AccessControlPolicy=value,
                    Bucket=bucket_name,
                    Key=key,
                )
            except ClientError as e:
                if e.response['Error']['Code'] == 'NoSuchKey':
                    logger.warning("No such key in S3: " + key)  # Will send the errors to the file
                else:
                    logger.error(("Unexpected error :{}".format(e)))
                    sys.exit(1)


def get_s3_acl(s3_bucket):
    s3_client = S3BotoWrapper(region_name=region)
    response_list = []
    try:
        s3_objects_key = get_all_s3_keys(s3_bucket, region)
    except ClientError as e:
        logger.error(("Unable to connect to AWS with error :{}".format(e)))
        sys.exit(1)
    for object_key in s3_objects_key:
        try:
            temp = {}
            response = s3_client.get_acl(Bucket=s3_bucket, Key=object_key)
            temp[object_key] = response
            response_list.append(temp)
        except ClientError as e:
            if e.response['Error']['Code'] == 'AccessDenied':
                logger.warning("You Don't have permission to access this object: " + object_key)
            elif e.response['Error']['Code'] == 'NoSuchKey':
                logger.warning("No such key in S3: " + object_key)  # Will send the errors to the file
            else:
                logger.error(("Unexpected error :{}".format(e)))
                sys.exit(1)
    return response_list


@click.command()
@click.option('--bucketname', required=True, help='S3 bucket name')
@click.option('--operation', required=True, help='Operation name to perform i.e 1- getacl 2- setaclprivate 3- revertacl')
def controller(bucketname, operation):
    file_to_write = bucketname + ".txt"
    if operation == 'getacl':
        objects_acl = get_s3_acl(bucketname)
        with open(file_to_write, 'w') as fout:
            json.dump(objects_acl, fout)
        logger.info("Task completed. Total numbers of objects read are: " + str(len(objects_acl)))
    elif operation == 'setaclprivate':
        try:
            data = []
            with open(file_to_write, "r") as inFile:
                data = json.load(inFile)
            set_acl_private(data, bucketname)
            logger.info("Task completed. ACL of " + bucketname + " objects set to private.")
        except IOError:
            logger.error("File not accessible")
            sys.exit(1)
    elif operation == 'revertacl':
        try:
            data = []
            with open(file_to_write, "r") as inFile:
                data = json.load(inFile)
            revert_s3_acl(data, bucketname)
            logger.info("Task completed. ACL of " + bucketname + " objects reverted to given state")
        except IOError:
            logger.error("File not accessible")
            sys.exit(1)
    else:
        logger.error("Invalid Operation. Please enter valid operation. Operation supported are i.e 1- getacl "
                     "2- setaclprivate 3- revertacl ")  # Will send the errors to the file
        sys.exit(0)


if __name__ == '__main__':
    controller()
<commit_msg>Updated script to take in a white list of prefixes to filte out<commit_after>#!/usr/bin/python3
"""
Get current ACL of all objects in given S3 bucket or set them to private or revert back.
Script supports 3 operations
1- getacl
2- setaclprivate
3- revertacl

1 optional parameter
whitelist (optional) (provide multiple whitelist parameters to filter out)

It saves current ACL in a file named bucketname.txt for updating or reverting purposes.

python s3_acl.py --bucketname <name-of-bucket> --operation getacl --whitelist <prefix_to_avoid>

Should assume role to run this script.
"""


import boto3
from botocore.exceptions import ClientError
import backoff
import sys
import json
import click
import logging

MAX_TRIES = 5
region = "us-east-1"
# Set logging configuration
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
# create file handler that logs messages
filehandler = logging.FileHandler('s3_acl.log')
filehandler.setLevel(logging.INFO)
formatter = logging.Formatter(
    '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
filehandler.setFormatter(formatter)
# add the handlers to logger
logger.addHandler(filehandler)


class S3BotoWrapper:
    def __init__(self, **kwargs):
        self.client = boto3.client("s3", **kwargs)

    @backoff.on_exception(backoff.expo,
                          ClientError,
                          max_tries=MAX_TRIES)
    def get_object(self, *args, **kwargs):
        return self.client.list_objects_v2(*args, **kwargs)

    @backoff.on_exception(backoff.expo,
                          ClientError,
                          max_tries=MAX_TRIES)
    def get_acl(self, *args, **kwargs):
        return self.client.get_object_acl(*args, **kwargs)

    @backoff.on_exception(backoff.expo,
                          ClientError,
                          max_tries=MAX_TRIES)
    def put_acl(self, *args, **kwargs):
        return self.client.put_object_acl(*args, **kwargs)


def get_all_s3_keys(s3_bucket, region, whitelist):
    """Get a list of all keys in an S3 bucket."""
    keys = []
    kwargs = {'Bucket': s3_bucket}
    while True:
        s3_client = S3BotoWrapper(region_name=region)
        resp = s3_client.get_object(**kwargs)
        for obj in resp['Contents']:
            # Filter out directories, you can add more filters here if required.
            if obj['Key'][-1] == '/' or any(obj['Key'].startswith(whitelist_object) for whitelist_object in whitelist):
                continue
            else:
                keys.append(obj['Key'])
        try:
            kwargs['ContinuationToken'] = resp['NextContinuationToken']
        except KeyError:
            break
    return keys


def set_acl_private(acl_list, bucket_name, whitelist):
    s3_client = S3BotoWrapper(region_name=region)
    for item in acl_list:
        for key, value in item.items():
            if any(key.startswith(whitelist_object) for whitelist_object in whitelist):
                continue
            else:
                try:
                    s3_client.put_acl(
                        ACL='private',
                        Bucket=bucket_name,
                        Key=key,
                    )
                except ClientError as e:
                    if e.response['Error']['Code'] == 'NoSuchKey':
                        logger.warning("No such key in S3: " + key)  # Will send the errors to the file
                    else:
                        logger.error(("Unexpected error :{}".format(e)))
                        sys.exit(1)


def revert_s3_acl(acl_list, bucket_name, whitelist):
    s3_client = S3BotoWrapper(region_name=region)
    for item in acl_list:
        for key, value in item.items():
            if any(key.startswith(whitelist_object) for whitelist_object in whitelist):
                continue
            else:
                try:
                    value.pop('ResponseMetadata', None)
                    s3_client.put_acl(
                        AccessControlPolicy=value,
                        Bucket=bucket_name,
                        Key=key,
                    )
                except ClientError as e:
                    if e.response['Error']['Code'] == 'NoSuchKey':
                        logger.warning("No such key in S3: " + key)  # Will send the errors to the file
                    else:
                        logger.error(("Unexpected error :{}".format(e)))
                        sys.exit(1)


def get_s3_acl(s3_bucket, whitelist):
    s3_client = S3BotoWrapper(region_name=region)
    response_list = []
    try:
        s3_objects_key = get_all_s3_keys(s3_bucket, region, whitelist)
    except ClientError as e:
        logger.error(("Unable to connect to AWS with error :{}".format(e)))
        sys.exit(1)
    for object_key in s3_objects_key:
        try:
            temp = {}
            response = s3_client.get_acl(Bucket=s3_bucket, Key=object_key)
            temp[object_key] = response
            response_list.append(temp)
        except ClientError as e:
            if e.response['Error']['Code'] == 'AccessDenied':
                logger.warning("You Don't have permission to access this object: " + object_key)
            elif e.response['Error']['Code'] == 'NoSuchKey':
                logger.warning("No such key in S3: " + object_key)  # Will send the errors to the file
            else:
                logger.error(("Unexpected error :{}".format(e)))
                sys.exit(1)
    return response_list


@click.command()
@click.option('--bucketname', required=True, help='S3 bucket name')
@click.option('--operation', required=True, help='Operation name to perform i.e 1- getacl 2- setaclprivate 3- revertacl')
@click.option('--whitelist', '-i', multiple=True, help='S3 objects name to avoid')
def controller(bucketname, operation, whitelist):
    file_to_write = bucketname + ".txt"
    if operation == 'getacl':
        objects_acl = get_s3_acl(bucketname, whitelist)
        with open(file_to_write, 'w') as fout:
            json.dump(objects_acl, fout)
        logger.info("Task completed. Total numbers of objects read are: " + str(len(objects_acl)))
    elif operation == 'setaclprivate':
        try:
            data = []
            with open(file_to_write, "r") as inFile:
                data = json.load(inFile)
            set_acl_private(data, bucketname, whitelist)
            logger.info("Task completed. ACL of " + bucketname + " objects set to private.")
        except IOError:
            logger.error("File not accessible")
            sys.exit(1)
    elif operation == 'revertacl':
        try:
            data = []
            with open(file_to_write, "r") as inFile:
                data = json.load(inFile)
            revert_s3_acl(data, bucketname, whitelist)
            logger.info("Task completed. ACL of " + bucketname + " objects reverted to given state")
        except IOError:
            logger.error("File not accessible")
            sys.exit(1)
    else:
        logger.error("Invalid Operation. Please enter valid operation. Operation supported are i.e 1- getacl "
                     "2- setaclprivate 3- revertacl ")  # Will send the errors to the file
        sys.exit(0)


if __name__ == '__main__':
    controller()

{'fraction_non_alphanumeric': 0.10333723653395785, 'fraction_numerical': 0.008782201405152224, 'mean_word_length': 2.1305841924398625, 'pattern_counts': {'":': 1, '<': 6, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5850760', 'n_tokens_mistral': 4025, 'n_tokens_neox': 3867, 'n_words': 1045}