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

Code before:
if config.get("development", False):
    # If debug mode, we will set up the schedule to run the next minute after the job starts
    time_to_start = datetime.utcnow() + timedelta(minutes=1)
    dev_schedule = crontab(hour=time_to_start.hour, minute=time_to_start.minute)
    schedule_30_minute = dev_schedule
    schedule_45_minute = dev_schedule
    schedule_6_hours = dev_schedule


@app.task(soft_time_limit=1800)
def get_inventory_of_iam_keys(account_id: str) -> bool:
    """
    This function will get all the AWS IAM Keys for all the IAM users in all the Netflix AWS accounts.

    - Create an Array of the AWS IAM Keys
    - Write to an S3 bucket : swag.mgmt.netflix.net/data/dirtylaundry.json
    - This will be picked up by DirtyLaundry
    """
    function: str = f"{__name__}.{sys._getframe().f_code.co_name}"
    # First, get list of accounts
    if config.region == config.get("celery.active_region"):
        key_data = []
        accounts_d: list = aws.get_account_ids_to_names()
        users: list = []
        for account_id in accounts_d.keys():
            iam_users = get_account_authorization_details(
                account_number=account_id,
                assume_role=config.get("policies.role_name"),
                region=config.region,
                filter="User",
            )

            for user in iam_users:
                kd = get_user_access_keys(account_number=account_id,
                                   assume_role=config.get("policies.role_name"),
                                   region=config.region, user=user)
                for key_details in kd:
                    key_data.append(key_details.get("AccessKeyId"))
        put_object(
            Bucket=config.get("get_inventory_of_iam_keys.Bucket"),
            assume_role=config.get("get_inventory_of_iam_keys.assume_role"),
            account_number=config.get("get_inventory_of_iam_keys.account_number"),
            region=config.get("get_inventory_of_iam_keys.region"),
            Key=config.get("get_inventory_of_iam_keys.Key"),
            Body=json.dumps(key_data),
            session_name='consoleme_dirtylaundry'
        )
    log_data = {
        "function": function,
        "message": "Successfully written to the S3 bucket"
    }

Review comment (bug):
What if this raises an exception because the ConsoleMe role isn't set up in one of the accounts that Swag returns? With this current code, it would raise an exception and the task would fail. Is this what we want? Or perhaps we can skip the account and write try/except with : log_data["message"] = "... Whatever the problem is .." , log_data["error"] = e, log.error(log_data, exc_info=True)

Code after:
    - Write this data to an S3 bucket
    """
    function: str = f"{__name__}.{sys._getframe().f_code.co_name}"
    log_data = {
        "function": function,
        "message": "Get inventory of IAM Keys"
    }
    # First, get list of accounts
    key_data = []
    if not config.get("get_inventory_of_iam_keys.enabled"):
        stats.count(f"{function}.success")
        return {}
    if config.region == config.get("celery.active_region") and config.get("environment") == "prod":
            accounts_d: list = aws.get_account_ids_to_names()
            users: list = []
            for account_id in accounts_d.keys():
                try:
                    iam_users = get_account_authorization_details(
                        account_number=account_id,
                        assume_role=config.get("policies.role_name"),
                        region=config.region,
                        filter="User",
                    )

                    for user in iam_users:
                        kd = get_user_access_keys(account_number=account_id,
                                                  assume_role=config.get("policies.role_name"),
                                                  region=config.region, user=user)
                        for key_details in kd:
                            key_data.append(key_details.get("AccessKeyId"))
                except Exception as e:
                    log_data["error"] = e
                    log.error(log_data, exc_info=True)
            put_object(
                Bucket=config.get("get_inventory_of_iam_keys.bucket"),
                assume_role=config.get("get_inventory_of_iam_keys.assume_role"),
                account_number=config.get("get_inventory_of_iam_keys.account_number"),
                region=config.get("get_inventory_of_iam_keys.region"),
                Key=config.get("get_inventory_of_iam_keys.key"),
                Body=json.dumps(key_data),
                session_name=config.get("get_inventory_of_iam_keys.session_name")
            )
    log_data["total_iam_access_key_id"] = len(key_data)
    log.debug(log_data)
    stats.count(f"{function}.success")
    return log_data


schedule_30_minute = timedelta(seconds=1800)
schedule_45_minute = timedelta(seconds=2700)
schedule_6_hours = timedelta(hours=6)