439 lines
15 KiB
Python
439 lines
15 KiB
Python
import dataclasses
|
|
import functools
|
|
import itertools
|
|
import logging
|
|
import re
|
|
import time
|
|
import typing
|
|
import uuid
|
|
from typing import Optional
|
|
|
|
import boto3
|
|
import botocore.exceptions
|
|
from botocore.exceptions import ClientError
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def retry(timeout: int, backoff_seconds: int = 10, bubble_errors: typing.List[typing.Type[Exception]] = None):
|
|
"""
|
|
Tries calling a function {timeout} seconds until it succeeds or gives up and throw an error
|
|
:param timeout: Timeout in seconds
|
|
:param backoff_seconds: Wait time between retries
|
|
:param bubble_errors: List of exceptions to bubble up
|
|
:return: Wrapped function
|
|
"""
|
|
bubble_errors = tuple(bubble_errors or [])
|
|
|
|
def wrapped(func: typing.Callable):
|
|
@functools.wraps(func)
|
|
def inner(*args, **kwargs):
|
|
waited = 0
|
|
while True:
|
|
try:
|
|
logger.debug(f"Calling {func}")
|
|
return func(*args, **kwargs)
|
|
except KeyboardInterrupt:
|
|
raise
|
|
except tuple(bubble_errors):
|
|
raise
|
|
except Exception as e:
|
|
doze = min(timeout - waited, backoff_seconds)
|
|
times = (timeout - waited) // backoff_seconds
|
|
logger.debug(f"{func} failed, will try {times} times in {doze} sec")
|
|
time.sleep(doze)
|
|
waited += backoff_seconds
|
|
if waited >= timeout:
|
|
raise TimeoutError(f"Couldn't get a successful result from {func} in {timeout} seconds") from e
|
|
|
|
return inner
|
|
|
|
return wrapped
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class RedisInstance:
|
|
cluster_id: str
|
|
host: str
|
|
port: int
|
|
arn: str
|
|
user_group_ids: typing.List[str] = dataclasses.field(default_factory=list)
|
|
|
|
|
|
def paginate_by_marker(
|
|
func: callable,
|
|
list_field: str,
|
|
) -> typing.Iterable:
|
|
"""
|
|
Simplifies paginating over AWS results with a simpler interface than Paginators.
|
|
:param func: A function that accepts a parameter named `Marker`
|
|
:param list_field: Result field to iterate on
|
|
"""
|
|
marker = ""
|
|
while True:
|
|
result = func(Marker=marker)
|
|
|
|
items = result.get(list_field, [])
|
|
yield from items
|
|
|
|
if "Marker" not in result:
|
|
break
|
|
|
|
marker = result["Marker"]
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class RedisUser:
|
|
username: str
|
|
password: str
|
|
redis_key_prefix: str
|
|
|
|
|
|
class CreateUserResult(typing.TypedDict):
|
|
user_group_id: str
|
|
users: typing.List[RedisUser]
|
|
status: str
|
|
|
|
|
|
class RedisService:
|
|
def __init__(self, region: str):
|
|
self.elasticache = boto3.client("elasticache", region_name=region)
|
|
|
|
def create_redis_acl_user(
|
|
self,
|
|
replication_group_id: str,
|
|
users: typing.List[RedisUser],
|
|
) -> None:
|
|
"""
|
|
Creates a Redis ACL user. It will throw an error if the user group attached to the replication group is not ready.
|
|
This prevents us from quickly creating users. In that case, pass in a list of users instead.
|
|
|
|
It takes about 2 minutes until the user is ready.
|
|
"""
|
|
repl_groups = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)
|
|
user_group_id: str = repl_groups['ReplicationGroups'][0]['UserGroupIds'][0]
|
|
group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)['UserGroups'][0]
|
|
|
|
"""
|
|
Limits imposed by AWS:
|
|
|
|
- User groups per replication group = 1
|
|
- Users per user group = 100
|
|
- Number of users = 1000
|
|
- Number of user groups = 100
|
|
https://docs.aws.amazon.com/AmazonElastiCache/latest/red-ug/Clusters.RBAC.html#Users-groups-to-RGs
|
|
|
|
In short, we can't place more than 100 users on a single Redis instance.
|
|
"""
|
|
|
|
total_users = len(group['UserIds'])
|
|
if total_users >= 100:
|
|
raise Exception('Replication group reached the limit of 100 users')
|
|
|
|
self._update_or_create_user_group(user_group_id, users)
|
|
|
|
def delete_redis_acl_user(self, username: str) -> None:
|
|
logger.info(f'Deleting user {username}')
|
|
self.elasticache.delete_user(UserId=username)
|
|
|
|
# Even though we just deleted the user, it takes a couple of seconds until AWS updates user group status.
|
|
# So, we have to wait a bit before we can actually wait and check that the changes have propagated.
|
|
# Overall, it should take about 2-3 minutes to delete a user.
|
|
|
|
def create_redis_acl(
|
|
self,
|
|
*,
|
|
replication_group_id: str,
|
|
description: str = None,
|
|
node_type: str,
|
|
cache_subnet_group_name: str,
|
|
security_group_id: str,
|
|
tags: typing.Dict[str, str] = None,
|
|
users: typing.List[RedisUser],
|
|
) -> dict:
|
|
user_group_id = f"{replication_group_id}-ug"
|
|
# we have to wait (~1m) until the user group is ready before we can create the ACL
|
|
self._update_or_create_user_group(user_group_id, users, wait=True)
|
|
|
|
cache_param_group = 'redis-acl-with-100-db'
|
|
self._ensure_cache_parameter_group(cache_param_group)
|
|
|
|
logger.info(f"Creating Redis replication group {replication_group_id=}")
|
|
default_replication_group_kwargs = dict(
|
|
ReplicasPerNodeGroup=0,
|
|
Engine="redis",
|
|
EngineVersion="6.x",
|
|
TransitEncryptionEnabled=True, # must be True to use Redis ACL
|
|
MultiAZEnabled=False,
|
|
AutomaticFailoverEnabled=False,
|
|
)
|
|
repl_result = self.elasticache.create_replication_group(
|
|
**default_replication_group_kwargs,
|
|
ReplicationGroupId=replication_group_id,
|
|
ReplicationGroupDescription=description or f'{replication_group_id} replication group',
|
|
CacheNodeType=node_type,
|
|
CacheSubnetGroupName=cache_subnet_group_name,
|
|
CacheParameterGroupName=cache_param_group,
|
|
UserGroupIds=[
|
|
user_group_id,
|
|
],
|
|
SecurityGroupIds=[
|
|
security_group_id,
|
|
],
|
|
Tags=[{"Key": k, "Value": v} for k, v in (tags or {}).items()],
|
|
)
|
|
|
|
# replication group takes ~8m minutes to be ready
|
|
|
|
return {
|
|
"replication_group_id": repl_result["ReplicationGroup"]["ReplicationGroupId"],
|
|
'arn': repl_result["ReplicationGroup"]["ARN"],
|
|
'status': repl_result["ReplicationGroup"]["Status"],
|
|
'users': users,
|
|
'_response': repl_result,
|
|
}
|
|
|
|
def _update_or_create_user_group(
|
|
self, user_group_id: str, users: typing.List[RedisUser], wait: bool = False
|
|
) -> None:
|
|
"""
|
|
Creates or updates a user group with the given users. Waits until the user group is active.
|
|
|
|
:param user_group_id:
|
|
:param users:
|
|
"""
|
|
# we have to add the `default` user for backwards compatibility
|
|
existing_users = self.elasticache.describe_users(
|
|
Filters=[{"Name": "user-id", "Values": [u.username for u in users]}],
|
|
)["Users"]
|
|
existing_user_ids = [it["UserId"] for it in existing_users]
|
|
|
|
created_user_ids = ["default"]
|
|
for user in users:
|
|
if user.username in existing_user_ids:
|
|
logger.info(f"User {user.username} is already exists")
|
|
created_user_ids.append(user.username)
|
|
continue
|
|
|
|
redis_acl = f"on +@all -@dangerous ~{user.redis_key_prefix}*"
|
|
logger.info(f"Creating tenant {user.username=}")
|
|
user_result = self.elasticache.create_user(
|
|
UserId=user.username,
|
|
UserName=user.username,
|
|
Passwords=[user.password],
|
|
AccessString=redis_acl,
|
|
Engine="redis",
|
|
)
|
|
created_user_ids.append(user_result["UserId"])
|
|
|
|
try:
|
|
group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)["UserGroups"][0]
|
|
member_user_ids: typing.List[str] = group["UserIds"]
|
|
users_to_add = set(created_user_ids) - set(member_user_ids)
|
|
|
|
if users_to_add:
|
|
self.elasticache.modify_user_group(
|
|
UserGroupId=user_group_id,
|
|
UserIdsToAdd=list(users_to_add),
|
|
)
|
|
# this will take some time (~45s) until the changes propagate
|
|
except self.elasticache.exceptions.UserGroupNotFoundFault:
|
|
logger.info('User group does not exist, creating')
|
|
_ = self.elasticache.create_user_group(
|
|
UserGroupId=user_group_id,
|
|
Engine="redis",
|
|
UserIds=created_user_ids,
|
|
)
|
|
# user group creation takes around 60s
|
|
if wait:
|
|
self._wait_user_group(user_group_id)
|
|
|
|
@retry(timeout=60, backoff_seconds=5)
|
|
def _wait_user_group(self, user_group_id: str) -> None:
|
|
logger.info(f'Checking status of user group {user_group_id=}')
|
|
g = self.elasticache.describe_user_groups(UserGroupId=user_group_id)["UserGroups"][0]
|
|
assert g["Status"] == "active"
|
|
|
|
def is_redis_acl_ready(self, replication_group_id: str) -> bool:
|
|
repl = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)[
|
|
'ReplicationGroups'
|
|
][0]
|
|
is_redis_available = repl['Status'] == 'available'
|
|
|
|
user_group_id = repl['UserGroupIds'][0]
|
|
group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)['UserGroups'][0]
|
|
is_group_available = group['Status'] == 'active'
|
|
|
|
return is_group_available and is_redis_available
|
|
|
|
def get_redis_acl(self, replication_group_id: str) -> dict:
|
|
redis = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)[
|
|
'ReplicationGroups'
|
|
][0]
|
|
return dict(
|
|
replication_group_id=redis['ReplicationGroupId'],
|
|
arn=redis['ARN'],
|
|
host=redis["NodeGroups"][0]["PrimaryEndpoint"]["Address"],
|
|
port=redis["NodeGroups"][0]["PrimaryEndpoint"]["Port"],
|
|
)
|
|
|
|
def delete_redis_acl(self, replication_group_id: str) -> None:
|
|
self.elasticache.delete_replication_group(ReplicationGroupId=replication_group_id)
|
|
|
|
def _ensure_cache_parameter_group(self, parameter_group_name: str) -> None:
|
|
try:
|
|
_ = self.elasticache.describe_cache_parameter_groups(CacheParameterGroupName=parameter_group_name)[
|
|
'CacheParameterGroups'
|
|
][0]
|
|
return
|
|
except self.elasticache.exceptions.CacheParameterGroupNotFoundFault:
|
|
pass
|
|
|
|
logger.info(f"Creating cache parameter group {parameter_group_name=}")
|
|
_ = self.elasticache.create_cache_parameter_group(
|
|
CacheParameterGroupName=parameter_group_name,
|
|
CacheParameterGroupFamily="redis6.x",
|
|
Description="Redis parameter group for the Redis ACL",
|
|
)
|
|
_ = self.elasticache.modify_cache_parameter_group(
|
|
CacheParameterGroupName=parameter_group_name,
|
|
ParameterNameValues=[
|
|
{
|
|
"ParameterName": "databases",
|
|
"ParameterValue": "100",
|
|
},
|
|
],
|
|
)
|
|
|
|
|
|
def akinon_create_redis(
|
|
*,
|
|
k8s_cluster_name: str,
|
|
region: str,
|
|
node_type: str,
|
|
replicas: int,
|
|
owner_arn: str,
|
|
app_name: str,
|
|
role: str,
|
|
):
|
|
subnet_group_name = f"{k8s_cluster_name}-redis-subg"
|
|
security_group_name = f"{k8s_cluster_name}-redis-sg"
|
|
|
|
# owner arn is formatted like:
|
|
# arn:aws:iam::412344683105:user/myusername
|
|
username = owner_arn.split("/")[1]
|
|
tags = {
|
|
"CostCenter": f"{username}-{app_name}-{role}-redis",
|
|
}
|
|
|
|
cache_cluster_id = uuid.uuid4().hex
|
|
r = RedisService(region)
|
|
return r.create_redis(
|
|
cluster_id=cache_cluster_id,
|
|
node_type=node_type,
|
|
replicas=replicas,
|
|
security_group_name=security_group_name,
|
|
subnet_group_name=subnet_group_name,
|
|
tags=tags,
|
|
)
|
|
|
|
|
|
def akinon_create_redis_with_acl(
|
|
k8s_cluster_name: str,
|
|
owner_arn: str,
|
|
node_type: str,
|
|
app_name: str,
|
|
role: str,
|
|
users: typing.List[RedisUser],
|
|
) -> RedisInstance:
|
|
cache_cluster_id = uuid.uuid4().hex
|
|
subnet_group_name = f"{k8s_cluster_name}-redis-subg"
|
|
security_group_name = f"{k8s_cluster_name}-redis-sg"
|
|
|
|
# owner arn is formatted as:
|
|
# arn:aws:iam::412344683105:user/myusername
|
|
username = owner_arn.split("/")[1]
|
|
tags = {
|
|
"CostCenter": f"{username}-{app_name}-{role}-redis",
|
|
}
|
|
|
|
return r.create_redis_acl(
|
|
cluster_id=cache_cluster_id,
|
|
node_type=node_type,
|
|
subnet_group_name=subnet_group_name,
|
|
security_group_name=security_group_name,
|
|
tags=tags,
|
|
users=users,
|
|
)
|
|
|
|
|
|
def chunk(it, size):
|
|
it = iter(it)
|
|
sentinel = ()
|
|
return iter(lambda: tuple(itertools.islice(it, size)), sentinel)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.INFO, format=f"%(asctime)s: {logging.BASIC_FORMAT}")
|
|
r = RedisService(region="eu-central-1")
|
|
# res = r.create_redis_acl(
|
|
# replication_group_id="zerodev3",
|
|
# description="zero redis dev",
|
|
# node_type="cache.t3.micro",
|
|
# cache_subnet_group_name="dev",
|
|
# security_group_id="sg-c11a77ac",
|
|
# tags={"CostCenter": "zero123"},
|
|
# users=[RedisUser(username="u1", password="zeropassword1234", redis_key_prefix="zero:")],
|
|
# )
|
|
# print(res)
|
|
|
|
r.create_redis_acl_user(
|
|
replication_group_id='zerodev3',
|
|
users=[
|
|
RedisUser(
|
|
username='u2',
|
|
password='testtesttesttest123',
|
|
redis_key_prefix='test:',
|
|
)
|
|
],
|
|
)
|
|
# r.delete_redis_acl_user('test6')
|
|
while True:
|
|
if r.is_redis_acl_ready('zerodev3'):
|
|
logger.info('ready')
|
|
print(r.get_redis_acl('zerodev3'))
|
|
break
|
|
logger.info('not ready')
|
|
time.sleep(10)
|
|
exit()
|
|
|
|
res = r.create_redis_acl(
|
|
replication_group_id="zerodev",
|
|
description="zero redis dev",
|
|
node_type="cache.t3.micro",
|
|
cache_subnet_group_name="dev",
|
|
security_group_id="sg-c11a77ac",
|
|
tags={"CostCenter": "zero123"},
|
|
users=[RedisUser(username="zerouser", password="zeropassword1234", redis_key_prefix="zero:")],
|
|
)
|
|
print(res)
|
|
#
|
|
# redis = r.get_redis_acl(cluster_arn="arn:aws:elasticache:eu-central-1:400344683105:replicationgroup:zero")
|
|
# r.add_redis_user(
|
|
# cluster_id=redis.cluster_id,
|
|
# users=[RedisUser(username="zero27", password="zero27password123", redis_key_prefix="zero27:")],
|
|
# )
|
|
# print(r.add_redis_user(cluster_id="zero1", users=[RedisUser("zero26", "zero26password123", "zero26:")]))
|
|
#
|
|
# print(
|
|
# r.create_redis_with_acl(
|
|
# cluster_id="zero",
|
|
# cluster_description="zero redis",
|
|
# node_type="cache.t3.micro",
|
|
# subnet_group_name="testing-subnet",
|
|
# security_group_name="default",
|
|
# tags={"CostCenter": "zero123"},
|
|
# users=[RedisUser(username="zerouser", password="zeropassword1234", redis_key_prefix="zero:")],
|
|
# )
|
|
# )
|