diff --git a/_variables.tf b/_variables.tf
index 87aa19f..9162a7a 100644
--- a/_variables.tf
+++ b/_variables.tf
@@ -90,6 +90,31 @@ variable "alb_sg_allow_test_listener" {
   description = "Whether to allow world access to the test listeners"
 }
 
+variable "alb_sg_custom_cidr_blocks" {
+  default     = ["0.0.0.0/0"]
+  description = "Block public access to ALB and stick to these CIDR blocks only"
+}
+
+variable "alb_additional_sg" {
+  default     = []
+  description = "pass addition list of security groups to add to ALB"
+}
+
+variable "alb_sg_allow_cloudfront" {
+  default     = false
+  description = "Whether to allow Cloudfront IP range to access"
+}
+
+variable "alb_sg_allow_api_gateway" {
+  default     = false
+  description = "Whether to allow API Gateway IP range to access"
+}
+
+variable "alb_sg_allow_api_gateway_region" {
+  default     = ""
+  description = "Use this region to allow API Gateway IP range"
+}
+
 variable "alb_sg_allow_egress_https_world" {
   default     = true
   description = "Whether to allow ALB to access HTTPS endpoints - needed when using OIDC authentication"
@@ -366,4 +391,16 @@ variable "container_insights" {
   type        = bool
   default     = false
   description = "Enables CloudWatch Container Insights for a cluster."
+}
+
+variable "enable_managed_draining" {
+  description = "Enable managed draining for ECS instances and add AmazonECSManaged tag"
+  type        = bool
+  default     = true
+}
+
+variable "enable_managed_termination_protection" {
+  description = "Enable managed draining for ECS instances and add AmazonECSManaged tag"
+  type        = bool
+  default     = false
 }
\ No newline at end of file
diff --git a/alb.tf b/alb.tf
index 54f353c..29067f1 100644
--- a/alb.tf
+++ b/alb.tf
@@ -8,9 +8,12 @@ resource "aws_lb" "ecs" {
   drop_invalid_header_fields = var.alb_drop_invalid_header_fields
   enable_deletion_protection = var.alb_enable_deletion_protection
 
-  security_groups = [
-    aws_security_group.alb[0].id,
-  ]
+  security_groups = compact(
+    concat(try(aws_security_group.from_cloudfront.*.id, []), 
+    try(var.alb_additional_sg,[]),
+    [ aws_security_group.alb[0].id,
+    try(aws_security_group.from_api_gateway[0].id, "")
+  ]))
 
   idle_timeout = 400
 
diff --git a/asg.tf b/asg.tf
index 787fbd9..759e5db 100644
--- a/asg.tf
+++ b/asg.tf
@@ -39,9 +39,19 @@ resource "aws_autoscaling_group" "ecs" {
     propagate_at_launch = true
   }
 
+  dynamic "tag" {
+    for_each = var.enable_managed_draining ? [1] : []
+    content {
+      key                 = "AmazonECSManaged"
+      value               = ""
+      propagate_at_launch = true
+    }
+  }
+
   target_group_arns         = var.target_group_arns
   health_check_grace_period = var.autoscaling_health_check_grace_period
   default_cooldown          = var.autoscaling_default_cooldown
+
   lifecycle {
     create_before_destroy = true
   }
@@ -53,7 +63,8 @@ resource "aws_ecs_capacity_provider" "ecs_capacity_provider" {
 
   auto_scaling_group_provider {
     auto_scaling_group_arn         = aws_autoscaling_group.ecs[0].arn
-    managed_termination_protection = "DISABLED"
+    managed_termination_protection = var.enable_managed_termination_protection ? "ENABLED" : "DISABLED"
+    managed_draining               = var.enable_managed_draining ? "ENABLED" : "DISABLED"
 
     managed_scaling {
       maximum_scaling_step_size = 10
@@ -62,4 +73,4 @@ resource "aws_ecs_capacity_provider" "ecs_capacity_provider" {
       target_capacity           = var.asg_target_capacity
     }
   }
-}
+}
\ No newline at end of file
diff --git a/sg-alb.tf b/sg-alb.tf
index b24554c..4e9828b 100644
--- a/sg-alb.tf
+++ b/sg-alb.tf
@@ -31,7 +31,7 @@ resource "aws_security_group_rule" "https_from_world_to_alb" {
   to_port           = 443
   protocol          = "tcp"
   security_group_id = aws_security_group.alb[0].id
-  cidr_blocks       = ["0.0.0.0/0"]
+  cidr_blocks       = var.alb_sg_custom_cidr_blocks
 }
 
 resource "aws_security_group_rule" "https_test_listener_from_world_to_alb" {
@@ -43,7 +43,7 @@ resource "aws_security_group_rule" "https_test_listener_from_world_to_alb" {
   to_port           = 8443
   protocol          = "tcp"
   security_group_id = aws_security_group.alb[0].id
-  cidr_blocks       = ["0.0.0.0/0"]
+  cidr_blocks       = var.alb_sg_custom_cidr_blocks
 }
 
 
diff --git a/sg-aws-iprange.tf b/sg-aws-iprange.tf
new file mode 100644
index 0000000..74665f6
--- /dev/null
+++ b/sg-aws-iprange.tf
@@ -0,0 +1,61 @@
+data "aws_ip_ranges" "cloudfront" {
+  regions  = ["global"]
+  services = ["cloudfront"]
+}
+
+data "aws_ip_ranges" "api_gateway" {
+  regions  = compact([data.aws_region.current.name,try(var.alb_sg_allow_api_gateway_region, "")])
+  services = ["api_gateway"]
+}
+
+locals {
+  ip_chunks = chunklist(data.aws_ip_ranges.cloudfront.cidr_blocks, 50)
+}
+
+resource "aws_security_group" "from_cloudfront" {
+  count = var.alb && var.alb_sg_allow_cloudfront ? length(local.ip_chunks) : 0
+  name = "from-cloudfront-${var.name}-${count.index}"
+  description = "SG for Request from Cloudfront"
+  vpc_id      = var.vpc_id
+
+  dynamic "ingress" {
+    for_each = local.ip_chunks[count.index]
+
+    content {
+      from_port   = 443
+      to_port     = 443
+      protocol    = "tcp"
+      cidr_blocks = [ingress.value]
+    }
+  }
+
+  tags = {
+    Name = "from-cloudfront-${var.name}-${count.index}"
+  }
+
+  lifecycle {
+    create_before_destroy = true
+  }
+}
+
+resource "aws_security_group" "from_api_gateway" {
+  count = var.alb && var.alb_sg_allow_api_gateway ? 1 : 0
+  name = "from-api-gateway-${var.name}"
+  description = "SG for Request from API Gateway"
+  vpc_id      = var.vpc_id
+
+  ingress {
+    from_port        = "443"
+    to_port          = "443"
+    protocol         = "tcp"
+    cidr_blocks      = data.aws_ip_ranges.api_gateway.cidr_blocks
+  }
+
+  tags = {
+    Name = "from-api-gateway-${var.name}"
+  }
+
+  lifecycle {
+    create_before_destroy = true
+  }
+}
\ No newline at end of file
diff --git a/userdata.tpl b/userdata.tpl
index 7cf538f..1c831a2 100644
--- a/userdata.tpl
+++ b/userdata.tpl
@@ -12,6 +12,20 @@ echo "vm.max_map_count=262144" >> /etc/sysctl.conf
 echo "fs.file-max=65536" >> /etc/sysctl.conf
 /sbin/sysctl -p /etc/sysctl.conf
 
+function checkECSregistration {
+    logger -s "ecs Check if there are any running Docker containers"
+    # Check if there are any running Docker containers
+    sleep 1200
+    docker_container_count=$(docker ps -q | wc -l)
+    ecs_registration=$(curl -s http://localhost:51678/v1/metadata | grep ${tf_cluster_name} | wc -l)
+    # If there are no running Docker containers and ECS agent is still not registered, terminate the instance
+    if [ "$docker_container_count" -le "1" ] && [ $ecs_registration -eq "0" ]; then
+        logger -s "No running Docker containers found and ECS agent is not registered. Terminating the instance..."
+        shutdown -h now
+    fi
+}
+
+checkECSregistration &
 
 echo "### INSTALL PACKAGES"
 yum update -y