#!/bin/python
import boto3
import sys
import os
import argparse
import logging
from multiprocessing import Process


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--accesskey', help='AWS Access key', default='')
    parser.add_argument('--secretkey', help='AWS Secret key', default='')
    parser.add_argument('--region', help='AWS Default region', default='us-east-1')
    parser.add_argument('--profile', help='AWS profile', default='')
    parser.add_argument('--concurrency', type=int, help='Size of a process pool', default=os.cpu_count())
    requiredNamed = parser.add_argument_group('required named arguments')
    requiredNamed.add_argument('--bucket', help='s3 bucket name', required=True)
    requiredNamed.add_argument('--canonicalId', help='Canonical ID of bucket owner user', required=True)
    requiredNamed.add_argument('--keyslistfile', help='file with list of keys', required=True)
    return parser.parse_args()

def set_object_bucket_owner_fc_acl(client,bucket,key,canonicalId):
    acl = client.get_object_acl(Bucket=bucket,Key=key)
    if (acl['Owner']['ID'] != canonicalId):
        return client.put_object_acl(ACL='bucket-owner-full-control',Bucket=bucket,Key=key)

def process_object_list(client,bucket,keylist,canonicalId):
    for key in keylist:
        set_object_bucket_owner_fc_acl(client,bucket,key,canonicalId)

def partitioned_keys_list(keys_list, partitions_num):
    chunk_size = int(len(keys_list)/partitions_num)
    for i in range(0, len(keys_list), chunk_size):
        if ((i + chunk_size) < len(keys_list)):
            yield keys_list[i:(i + chunk_size)]
        else:
            yield keys_list[i:]

if __name__ == "__main__":
    logger = logging.getLogger('')
    logger.setLevel(logging.DEBUG)

    args = parse_arguments()

    # Initialize the client with credentials or with IAM role.
    if (args.accesskey and args.secretkey):
        s3_client = boto3.client(
          's3',
          aws_access_key_id = args.accesskey,
          aws_secret_access_key = args.secretkey
        )
    elif args.profile:
        session = boto3.Session(profile_name=args.profile)
        s3_client = session.client('s3')
    else:
        s3_client = boto3.client('s3')

    # I used aws cli for creating keys list
    # aws s3api list-objects --profile debug_profile --bucket bucket_name --query "(Contents)[].[Key]" --output text
    keys_list = [line.rstrip('\n') for line in open(args.keyslistfile, 'r')]
    for partitioned_list in partitioned_keys_list(keys_list, args.concurrency):
        print(len(partitioned_list))
        p = Process(target=process_object_list, args=(s3_client,args.bucket,partitioned_list,args.canonicalId))
        p.start()