Repository: Netflix/consoleme
File: consoleme/celery/celery_tasks.py (Python)

Code before:
@app.task(soft_time_limit=1800)
def get_iam_role_limit() -> bool:
    """
    This function will gather the number of existing IAM Roles and IAM Role quota in all the AWS accounts.
    """

    function: str = f"{__name__}.{sys._getframe().f_code.co_name}"
    # First, get list of accounts

    if not config.get("get_iam_role_limit.enabled"):
        stats.count(f"{function}.success")
        return {}
    if (
            config.region == config.get("celery.active_region")
            and config.get("environment") == "prod"
    ):

        @sts_conn('iam')
        def _get_delivery_channels(**kwargs) -> list:
            """Gets the delivery channels in the account/region -- calls are wrapped with CloudAux"""
            return kwargs.pop("client").get_account_summary(**kwargs)

        if config.region == config.get("celery.active_region"):
            accounts_d: list = aws.get_account_ids_to_names()
            for account_id in accounts_d.keys():
                iam_summary = _get_delivery_channels(
                    account_number=account_id,
                    assume_role=config.get("policies.role_name"),
                    region=config.region,
                )
                log_data = {
                    "function": function,
                    "message": "Total number of IAM roles and IAM roles quota.",
                    "iam_roles": iam_summary["SummaryMap"]["Roles"],
                    "iam_quota": iam_summary["SummaryMap"]["RolesQuota"],
                    "aws_account": account_id,
                }
                stats.count(f"{function}.success", tags={"total_iam_roles": iam_summary["SummaryMap"]["Roles"],
                                                         "role_quota": iam_summary["SummaryMap"]["RolesQuota"],
                                                         "account_id": account_id})
                log.debug(log_data)

    log_data = {
        "function": function,
        "message": "Completed : looking for IAM Role limits",
    }
    log.debug(log_data)
    stats.count(f"{function}.success")
    return True



Review comment (bug):
What if we fail to assume the role in the account? We should have logic that will catch and report any exceptions, but allow the other accounts to complete successfully

Code after:
@app.task(soft_time_limit=1800)
def get_iam_role_limit() -> dict:
    """
    This function will gather the number of existing IAM Roles and IAM Role quota in all owned AWS accounts.
    """

    function: str = f"{__name__}.{sys._getframe().f_code.co_name}"

    if not config.get("celery.get_iam_role_limit.enabled"):
        return {}
    if config.region == config.get("celery.active_region") and config.get(
        "environment"
    ) in ["prod", "dev"]:

        @sts_conn("iam")
        def _get_delivery_channels(**kwargs) -> list:
            """Gets the delivery channels in the account/region -- calls are wrapped with CloudAux"""
            return kwargs.pop("client").get_account_summary(**kwargs)

        if config.region == config.get("celery.active_region"):
            # First, get list of accounts
            accounts_d: list = aws.get_account_ids_to_names()
            num_accounts = len(accounts_d.keys())
            num_roles = 0
            for account_id in accounts_d.keys():
                try:
                    iam_summary = _get_delivery_channels(
                        account_number=account_id,
                        assume_role=config.get("policies.role_name"),
                        region=config.region,
                    )
                    num_iam_roles = iam_summary["SummaryMap"]["Roles"]
                    iam_role_quota = iam_summary["SummaryMap"]["RolesQuota"]
                    iam_role_quota_ratio = num_iam_roles / iam_role_quota

                    num_roles += num_iam_roles
                    log_data = {
                        "function": function,
                        "message": "IAM role quota for account",
                        "num_iam_roles": num_iam_roles,
                        "iam_role_quota": iam_role_quota,
                        "iam_role_quota_ratio": iam_role_quota_ratio,
                        "account_id": account_id,
                    }
                    stats.count(
                        f"{function}.quota_ratio",
                        tags={
                            "num_iam_roles": num_iam_roles,
                            "iam_role_quota": iam_role_quota,
                            "iam_role_quota_ratio": iam_role_quota_ratio,
                            "account_id": account_id,