def update_policy()

in functions/source/c1c_controltower_lifecycle.py [0:0]

• 48 lines of code
• 6 McCabe index (conditional complexity)

def update_policy(aws_account_id):
    logger.info(f"Updating account {aws_account_id}")
    c1c_policy_document = c1cresources.ConformityPolicyDoc()
    sts_session = assume_role(aws_account_id, c1cresources.ControlTowerRoleName)

    sts_client = sts_session.client("sts")
    sts_identity = sts_client.get_caller_identity()
    partition = sts_identity["Arn"].split(":")[1]

    client = sts_session.client("iam")
    policy_resource = sts_session.resource("iam")
    logger.info(f"Updating policy in account {aws_account_id}")
    try:
        client.get_role(RoleName=c1cresources.IamRoleName)
    except client.exceptions.NoSuchEntityException:
        logger.info(f"Policy not found; configuring account")
        configure_account(aws_account_id)
        return

    logger.info(f"Updating AssumeRolePolicyDocument in account {aws_account_id}")
    try:
        c1c_connector = c1cconnectorapi.CloudOneConformityConnector(
            c1cresources.get_api_key()
        )
        client.update_assume_role_policy(
            RoleName=c1cresources.IamRoleName,
            PolicyDocument=c1cresources.get_assume_role_policy_document(c1c_connector),
        )
    except Exception as e:
        logger.error(f"Failed to update AssumeRolePolicyDocument: {e}")
        raise
    try:
        policy_part = 0
        # TODO this won't work if the partition is not aws.
        for policy in c1c_policy_document.list_of_policies:
            current_policy_object = policy_resource.Policy(
                f"arn:{partition}:iam::{aws_account_id}:policy/{c1cresources.IamPolicyName}{policy_part}"
            )
            current_policy_object_version = current_policy_object.default_version
            client.create_policy_version(
                PolicyArn=f"arn:{partition}:iam::{aws_account_id}:policy/{c1cresources.IamPolicyName}{policy_part}",
                PolicyDocument=json.dumps(policy.get("document")),
                SetAsDefault=True,
            )
            client.delete_policy_version(
                PolicyArn=f"arn:{partition}:iam::{aws_account_id}:policy/{c1cresources.IamPolicyName}{policy_part}",
                VersionId=current_policy_object_version.version_id,
            )
            policy_part += 1
    except Exception as e:
        logger.error(f"Failed to update policy {e}")
        raise