from xmlrpc.client import Boolean
import boto3
import logging
from datetime import date, datetime
from botocore.exceptions import ClientError
import json
from json_datetime_serializer import json_datetime_serializer
from kms_client import kmsClient
from search_kms_using_account_id import search_kms_using_account_id
from get_arn import get_arn

logger = logging.getLogger()
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s: %(levelname)s: %(message)s')


# mrk-8b618d3b5c1d432aad977f31b29df64b
# 862638055046

def update_kms_key_policy(target_arn, key_id):
    client = kmsClient()
    kms_policy_data = client.get_key_policy(KeyId=key_id,PolicyName='default')
    kms_policy = json.loads(kms_policy_data["Policy"])
    policy_update = Boolean(False)
    for statement in kms_policy["Statement"]:
        if "Allow use of the key" in statement["Sid"]:
            if target_arn not in statement['Principal']['AWS']:
                if type(statement['Principal']['AWS']) is list:
                    statement['Principal']['AWS'].append(target_arn)
                    statement['Principal']['AWS'].sort()
                    policy_update = True
            else:
                print('Target ARN already exists in the policy')
                policy_update = False
    if policy_update:
        response = client.put_key_policy(
            KeyId=key_id,
            PolicyName='default',
            Policy=json.dumps(kms_policy)
        )
        print(response)
    return kms_policy


if __name__ == '__main__':
    
    ARN = 'arn:aws:iam::862638055046:user/test-user'
    print(f"Arn: {ARN}")
    kms = update_kms_key_policy(ARN)
    logger.info('Getting information about KMS key...')
    logger.info(
        f'Key Details: {json.dumps(kms, indent=4, default=json_datetime_serializer)}'
    )
    logger.info('Done!')

from xmlrpc.client import Boolean
import boto3
import logging
from datetime import date, datetime
from botocore.exceptions import ClientError
import json
from json_datetime_serializer import json_datetime_serializer
from kms_client import kmsClient
from search_kms_using_account_id import search_kms_using_account_id
from get_arn import get_arn

logger = logging.getLogger()
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s: %(levelname)s: %(message)s')


# mrk-8b618d3b5c1d432aad977f31b29df64b
# 862638055046

def update_kms_key_policy(target_arn, key_id):
    client = kmsClient()
    kms_policy_data = client.get_key_policy(KeyId=key_id,PolicyName='default')
    kms_policy = json.loads(kms_policy_data["Policy"])
    policy_update = Boolean(False)
    for statement in kms_policy["Statement"]:
        if "Allow use of the key" in statement["Sid"]:
            if target_arn not in statement['Principal']['AWS']:
                if type(statement['Principal']['AWS']) is list:
                    statement['Principal']['AWS'].append(target_arn)
                    statement['Principal']['AWS'].sort()
                    policy_update = True
            else:
                print('Target ARN already exists in the policy')
                policy_update = False
    if policy_update:
        response = client.put_key_policy(
            KeyId=key_id,
            PolicyName='default',
            Policy=json.dumps(kms_policy)
        )
        print(response)
    return kms_policy


if __name__ == '__main__':
    
    ARN = 'arn:aws:iam::862638055046:user/test-user'
    print(f"Arn: {ARN}")
    kms = update_kms_key_policy(ARN)
    logger.info('Getting information about KMS key...')
    logger.info(
        f'Key Details: {json.dumps(kms, indent=4, default=json_datetime_serializer)}'
    )
    logger.info('Done!')

# mrk-8b618d3b5c1d432aad977f31b29df64b # 862638055046