def test_clear_old_redis_iam_cache(self):
from consoleme.config.config import CONFIG
from consoleme.lib.redis import RedisHandler
red = RedisHandler().redis_sync()
self.celery.REDIS_IAM_COUNT = 3
# Clear out the existing cache from Redis:
red.delete("test_cache_roles_for_account_expiration")
# Set the config value for the redis cache location
old_value = CONFIG.config["aws"].pop("iamroles_redis_key", None)
CONFIG.config["aws"][
"iamroles_redis_key"
] = "test_cache_roles_for_account_expiration"
# Add in some dummy IAM roles with a TTL that is more than 6 hours old:
old_ttl = int((datetime.utcnow() - timedelta(hours=6, seconds=5)).timestamp())
# 13 items / 3 = 5 iterations -- all of these roles should be cleaned up:
for i in range(0, 13):
role_entry = {
"arn": f"arn:aws:iam::123456789012:role/RoleNumber{i}",
"name": f"RoleNumber{i}",
"accountId": "123456789012",
"ttl": old_ttl,
"policy": "{}",
}
self.celery._add_role_to_redis(
"test_cache_roles_for_account_expiration", role_entry
)
# Add a role with a current TTL -- this should not be cleaned up:
role_entry = {
"arn": "arn:aws:iam::123456789012:role/RoleNumber99",
"name": "RoleNumber99",
"accountId": "123456789012",
"ttl": int(datetime.utcnow().timestamp()),
"policy": "{}",
}
self.celery._add_role_to_redis(
"test_cache_roles_for_account_expiration", role_entry
)
# Nothing should happen if we are not in us-west-2:
old_conf_region = self.celery.config.region
self.celery.config.region = "us-east-1"
self.celery.clear_old_redis_iam_cache()
self.assertEqual(red.hlen("test_cache_roles_for_account_expiration"), 14)
# With the proper region:
self.celery.config.region = "us-west-2"
self.celery.clear_old_redis_iam_cache()
# Verify:
self.assertEqual(red.hlen("test_cache_roles_for_account_expiration"), 1)
self.assertIsNotNone(
red.hget(
"test_cache_roles_for_account_expiration",
"arn:aws:iam::123456789012:role/RoleNumber99",
)
)
# Clear out the existing cache from Redis:
red.delete("test_cache_roles_for_account_expiration")
# Reset the config values:
self.celery.config.region = old_conf_region
self.celery.REDIS_IAM_COUNT = 1000
if not old_value:
del CONFIG.config["aws"]["iamroles_redis_key"]
else:
CONFIG.config["aws"]["iamroles_redis_key"] = old_value
def test_cache_roles_for_account(self):
from consoleme.config.config import CONFIG
from consoleme.lib.dynamo import IAMRoleDynamoHandler
from consoleme.lib.redis import RedisHandler
red = RedisHandler().redis_sync()
# Set the config value for the redis cache location
old_value = CONFIG.config["aws"].pop("iamroles_redis_key", None)
CONFIG.config["aws"]["iamroles_redis_key"] = "test_cache_roles_for_account"
# Clear out the existing cache from Redis:
red.delete("test_cache_roles_for_account")
# Run it:
self.celery.cache_roles_for_account("123456789012")
# Verify that everything is there:
dynamo = IAMRoleDynamoHandler()
results = dynamo.role_table.scan(TableName="consoleme_iamroles_global")
remaining_roles = [
"arn:aws:iam::123456789012:role/ConsoleMe",
"arn:aws:iam::123456789012:role/cm_someuser_N",
"arn:aws:iam::123456789012:role/awsaccount_user",
"arn:aws:iam::123456789012:role/TestInstanceProfile",
"arn:aws:iam::123456789012:role/rolename",
] + [f"arn:aws:iam::123456789012:role/RoleNumber{num}" for num in range(0, 10)]
self.assertEqual(results["Count"], len(remaining_roles))
self.assertEqual(results["Count"], red.hlen("test_cache_roles_for_account"))
for i in results["Items"]:
remaining_roles.remove(i["arn"])
self.assertEqual(i["accountId"], "123456789012")
self.assertGreater(int(i["ttl"]), 0)
self.assertIsNotNone(json.loads(i["policy"]))
self.assertEqual(
json.loads(red.hget("test_cache_roles_for_account", i["arn"]))[
"policy"
],
i["policy"],
)
# Should all be accounted for:
self.assertEqual(remaining_roles, [])
# We should have the same data in redis on all regions, this time coming from DDB
old_conf_region = self.celery.config.region
self.celery.config.region = "us-east-1"
# Clear out the existing cache from Redis:
red.delete("test_cache_roles_for_account")
# This should spin off extra fake celery tasks
res = self.celery.cache_roles_across_accounts()
self.assertEqual(
res,
{
"function": "consoleme.celery_tasks.celery_tasks.cache_roles_across_accounts",
"cache_key": "test_cache_roles_for_account",
"num_roles": 0,
"num_accounts": 1,
},
)
# Reset the config value:
self.celery.config.region = old_conf_region
if not old_value:
del CONFIG.config["aws"]["iamroles_redis_key"]
else:
CONFIG.config["aws"]["iamroles_redis_key"] = old_value
