initial commit
This commit is contained in:
@@ -0,0 +1,438 @@
|
||||
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:")],
|
||||
# )
|
||||
# )
|
||||
Reference in New Issue
Block a user