def wait_for_fargate_cluster_status(
cluster_name: str,
cluster_stats: ClusterStatus,
timeout_seconds: int = CLUSTER_PROVISION_TIMEOUT,
) -> None:
"""
Waits for a cluster to to reach a desired status by polling the current
state of the cluster
"""
cluster_provisioned = False
wait_time = 0
log_info(f"Waiting for cluster to reach {cluster_stats.value} state...")
while not cluster_provisioned and wait_time < timeout_seconds:
cluster_info = describe_fargate_cluster(cluster_name)
if len(cluster_info["failures"]) > 0:
break
cluster_provisioned = all([
c["status"] == cluster_stats.value
for c in cluster_info["clusters"]
], )
sleep(2)
wait_time += 2
if not cluster_provisioned:
log_error("Cluster failed to provision")
raise Abort()
def wait_for_tasks_to_start(
cluster_name: str,
tasks: List[TaskTypeDef],
timeout_seconds: int = TASK_BOOT_TIMEOUT,
) -> None:
"""
Waits for all of the tasks to reach their desired state by polling
the current state of the tasks
"""
task_arns = [t["taskArn"] for t in tasks]
tasks_started = False
wait_time = 0
log_info("Waiting for bastion task to start...")
while not tasks_started and wait_time < timeout_seconds:
task_info = describe_task(cluster_name, task_arns)
if not task_info or len(task_info["failures"]) > 0:
break
tasks_started = all([
t["lastStatus"] == t["desiredStatus"] for t in task_info["tasks"]
], )
sleep(2)
wait_time += 2
if not tasks_started:
log_error("Bastion task failed to start")
raise Abort()
def delete_task_definition() -> None:
"""
Inactivates all serverless-aws-bastion task definitions
"""
client: ECSClient = fetch_boto3_client("ecs")
task_definitions = client.list_task_definitions(familyPrefix=DEFAULT_NAME)
log_info("Deregistering task definitions")
for td in task_definitions["taskDefinitionArns"]:
client.deregister_task_definition(taskDefinition=td)
def delete_role(role_name: str) -> None:
"""
Safely deletes a given role by first detaching any policies
and then deleting the role and handling any exceptions
"""
client: IAMClient = fetch_boto3_client("iam")
try:
log_info(f"Deleting {role_name} role")
detach_policies_from_role(role_name)
client.delete_role(RoleName=role_name)
except client.exceptions.NoSuchEntityException:
return None
def create_fargate_cluster(cluster_name: str) -> CreateClusterResponseTypeDef:
"""
Creates a Fargate cluster to launch the bastion task into
"""
client: ECSClient = fetch_boto3_client("ecs")
log_info("Creating Fargate cluster")
response = client.create_cluster(
clusterName=cluster_name,
capacityProviders=["FARGATE"],
tags=build_tags("ecs"),
)
wait_for_fargate_cluster_status(cluster_name, ClusterStatus.ACTIVE)
return response
def delete_deregister_ssm_policy() -> None:
"""
Deletes the IAM policy that allows the bastion ECS task to
deregister itself from SSM.
"""
client: IAMClient = fetch_boto3_client("iam")
try:
log_info(f"Deleting {SSM_DEREGISTER_POLICY_NAME} policy")
account_id = load_aws_account_id()
client.delete_policy(
PolicyArn=
f"arn:aws:iam::{account_id}:policy/{SSM_DEREGISTER_POLICY_NAME}", )
except client.exceptions.NoSuchEntityException:
return None
def delete_fargate_cluster(cluster_name: str) -> None:
"""
Deletes a given Fargate cluster
"""
client: ECSClient = fetch_boto3_client("ecs")
log_info("Deleting Fargate cluster")
try:
client.delete_cluster(cluster=cluster_name)
except client.exceptions.ClusterNotFoundException:
log_error(f"Failed to find {cluster_name} Fargate cluster")
raise Abort()
wait_for_fargate_cluster_status(cluster_name, ClusterStatus.INACTIVE)
def create_bastion_task_role() -> str:
"""
Creates the role that will be used by the bastion ECS task.
Skips creation if the role already exists.
Returns role arn
"""
client: IAMClient = fetch_boto3_client("iam")
current_role_arn = fetch_role_arn(TASK_ROLE_NAME)
if current_role_arn:
return current_role_arn
log_info(f"Creating {TASK_ROLE_NAME} role")
response = client.create_role(
RoleName=TASK_ROLE_NAME,
Description="Used by serverless-aws-bastion ECS tasks",
AssumeRolePolicyDocument=json.dumps(
{
"Version":
"2012-10-17",
"Statement": [
{
"Sid": "",
"Effect": "Allow",
"Principal": {
"Service":
["ecs-tasks.amazonaws.com", "ssm.amazonaws.com"],
},
"Action": "sts:AssumeRole",
},
],
}, ),
Tags=build_tags("iam"),
)
deregister_ssm_arn = create_deregister_ssm_policy()
attach_policies_to_role(
TASK_ROLE_NAME,
[
deregister_ssm_arn,
"arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore",
],
)
return response["Role"]["Arn"]
def create_bastion_task_execution_role() -> str:
"""
Creates the role that will be used by ECS to launch the bastion ECS task.
Skips creation if the role already exists.
Returns role arn
"""
client: IAMClient = fetch_boto3_client("iam")
current_role_arn = fetch_role_arn(TASK_EXECUTION_ROLE_NAME)
if current_role_arn:
return current_role_arn
log_info(f"Creating {TASK_EXECUTION_ROLE_NAME} role")
response = client.create_role(
RoleName=TASK_EXECUTION_ROLE_NAME,
Description=
"Used by Fargate to launch serverless-aws-bastion ECS tasks",
AssumeRolePolicyDocument=json.dumps(
{
"Version":
"2012-10-17",
"Statement": [
{
"Sid": "",
"Effect": "Allow",
"Principal": {
"Service": ["ecs-tasks.amazonaws.com"]
},
"Action": "sts:AssumeRole",
},
],
}, ),
Tags=build_tags("iam"),
)
attach_policies_to_role(
TASK_EXECUTION_ROLE_NAME,
[
"arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy",
],
)
return response["Role"]["Arn"]
def create_task_definition(task_role_arn: str,
execution_role_arn: str) -> None:
"""
Creates the task definition that will be used to launch the
serverless bastion container
"""
client: ECSClient = fetch_boto3_client("ecs")
log_info("Creating bastion ECS task")
client.register_task_definition(
family=DEFAULT_NAME,
networkMode="awsvpc",
cpu=TASK_CPU,
memory=TASK_MEMORY,
taskRoleArn=task_role_arn,
executionRoleArn=execution_role_arn,
containerDefinitions=[
{
"image":
f"nplutt/{DEFAULT_NAME}",
"name":
DEFAULT_NAME,
"essential":
True,
"portMappings": [
{
"hostPort": 22,
"protocol": "tcp",
"containerPort": 22,
},
],
"logConfiguration": {
"logDriver": "awslogs",
"options": {
"awslogs-group": "/ecs/ssh-bastion",
"awslogs-region": load_aws_region_name(),
"awslogs-stream-prefix": "ecs",
},
},
},
],
tags=build_tags("ecs"),
)
def create_deregister_ssm_policy() -> str:
"""
Creates an IAM policy that allows the bastion ECS task to
deregister itself from SSM.
Returns the policy arn
"""
client: IAMClient = fetch_boto3_client("iam")
try:
log_info(f"Creating {SSM_DEREGISTER_POLICY_NAME} policy")
response = client.create_policy(
Description="Used by serverless-aws-bastion ECS task to "
"deregister itself from SSM",
PolicyName=SSM_DEREGISTER_POLICY_NAME,
PolicyDocument=json.dumps(
{
"Version":
"2012-10-17",
"Statement": [
{
"Action": [
"ssm:DeregisterManagedInstance",
"ssm:DescribeInstanceInformation",
],
"Effect":
"Allow",
"Resource":
"*",
},
],
}, ),
)
return response["Policy"]["Arn"]
except client.exceptions.EntityAlreadyExistsException:
account_id = load_aws_account_id()
return f"arn:aws:iam::{account_id}:policy/{SSM_DEREGISTER_POLICY_NAME}"
def stop_fargate_tasks(cluster: str, tasks: List[TaskTypeDef]) -> None:
client: ECSClient = fetch_boto3_client("ecs")
log_info(f"Stopping {len(tasks)} tasks...")
for t in tasks:
client.stop_task(cluster=cluster, task=t["taskArn"])
def launch_fargate_task(
cluster_name: str,
subnet_ids: str,
security_group_ids: str,
authorized_keys: str,
instance_name: str,
timeout_minutes: int,
bastion_type: BastionType,
) -> RunTaskResponseTypeDef:
"""
Launches the ssh bastion Fargate task into the proper subnets & security groups,
also sends in the authorized keys.
"""
client: ECSClient = fetch_boto3_client("ecs")
bastion_id = str(uuid4())
activation: Dict[str, str] = {}
if bastion_type == BastionType.ssm:
activation = create_activation(TASK_ROLE_NAME, instance_name,
bastion_id) # type: ignore
log_info("Starting bastion task")
try:
response = client.run_task(
cluster=cluster_name,
taskDefinition=DEFAULT_NAME,
overrides={
"containerOverrides": [
{
"name":
DEFAULT_NAME,
"environment": [
{
"name": "AUTHORIZED_SSH_KEYS",
"value": authorized_keys
},
{
"name": "ACTIVATION_ID",
"value": activation.get("ActivationId", ""),
},
{
"name": "ACTIVATION_CODE",
"value": activation.get("ActivationCode", ""),
},
{
"name": "AWS_REGION",
"value": load_aws_region_name()
},
{
"name": "TIMEOUT",
"value": str(timeout_minutes * 60)
},
{
"name": "BASTION_TYPE",
"value": bastion_type.value
},
],
},
],
},
count=1,
launchType="FARGATE",
networkConfiguration={
"awsvpcConfiguration": {
"subnets": subnet_ids.split(","),
"securityGroups": security_group_ids.split(","),
"assignPublicIp": "ENABLED",
},
},
tags=build_tags(
"ecs",
{
"Name": f"{DEFAULT_NAME}/{instance_name}",
"BastionId": bastion_id,
"ActivationId": activation.get("ActivationId", ""),
},
),
)
except client.exceptions.ClusterNotFoundException:
log_error(
"Specified cluster to launch bastion task into doesn't exist")
raise Abort()
except (
client.exceptions.ClientException,
client.exceptions.InvalidParameterException,
) as e:
log_error(e.response["Error"]["Message"])
raise Abort()
wait_for_tasks_to_start(cluster_name, response["tasks"])
return response