fix: add support for multiple azs (closes #17) (#18)

* add support for multiple azs

* update tests

* remove jfiola from defaults

* update asg lambda role policy to include all management subnet arns

---------

Co-authored-by: Jacob Fiola <[email protected]>

Author: jmfiola

README.md

@@ -8,7 +8,8 @@ Terraform for Corelight's AWS Cloud Sensor Deployment.
 ```terraform
 
 data "aws_subnet" "management" {
-  id = "<management subnet id>"
+  for_each = toset(["<management subnet id 1>", "<management subnet id 2>"])
+  id       = each.value
 }
 
 module "asg_lambda_role" {
@@ -17,26 +18,31 @@ module "asg_lambda_role" {
   lambda_cloudwatch_log_group_arn = module.sensor.cloudwatch_log_group_arn
   sensor_autoscaling_group_arn    = module.sensor.autoscaling_group_arn
   security_group_arn              = module.sensor.management_security_group_arn
-  subnet_arn                      = data.aws_subnet.management.arn
+  subnet_arns                     = [for subnet in data.aws_subnet.management : subnet.arn]
 }
 
 module "sensor" {
   source = "github.com/corelight/terraform-aws-sensor"
 
-  # Recommend deploying a sensor per availability zone. Multiple AZs can
-  # be set but GWLB cross availability zone support is not recommended.
-  auto_scaling_availability_zones = ["<availability zone>"]
+  # Multi-AZ support: provide one subnet per availability zone
+  # The ASG will automatically distribute instances across AZs
+  availability_zones = ["us-east-1a", "us-east-1b"]
   aws_key_pair_name = "<key pair name>"
 
   # Request access to Corelight sensor AMI from you Account Executive
   corelight_sensor_ami_id = "<sensor AMI ID>"
   license_key = "<your Corelight sensor license key>"
-  management_subnet_id = "<management subnet>"
-  monitoring_subnet_id = "<monitoring subnet>"
+
+  # Provide one management subnet per AZ (must match availability_zones order)
+  management_subnet_ids = ["<management subnet in us-east-1a>", "<management subnet in us-east-1b>"]
+
+  # Provide one monitoring subnet per AZ (must match availability_zones order)
+  monitoring_subnet_ids = ["<monitoring subnet in us-east-1a>", "<monitoring subnet in us-east-1b>"]
+
   community_string = "<password for the sensor api>"
   vpc_id = "<vpc where the sensor autoscaling group is deployed>"
   asg_lambda_iam_role_arn = module.asg_lambda_role.role_arn
-  
+
   fleet_token = "<the pairing token from the Fleet UI>"
   fleet_url   = "<the URL of the fleet instance from the Fleet UI>"
   fleet_server_sslname = "<the ssl name provided by Fleet>"

autoscaling_group.tf

@@ -9,7 +9,7 @@ resource "aws_autoscaling_group" "sensor_asg" {
     version = aws_launch_template.sensor_launch_template.latest_version
   }
 
-  availability_zones        = var.availability_zones
+  vpc_zone_identifier       = var.monitoring_subnet_ids
   target_group_arns         = [aws_lb_target_group.health_check.arn]
   health_check_type         = "EC2"
   health_check_grace_period = 300

data.tf

@@ -2,6 +2,12 @@ data "aws_vpc" "provided" {
   id = var.vpc_id
 }
 
-data "aws_subnet" "monitoring_subnet" {
-  id = var.monitoring_subnet_id
+data "aws_subnet" "monitoring_subnets" {
+  for_each = toset(var.monitoring_subnet_ids)
+  id       = each.value
+}
+
+data "aws_subnet" "management_subnets" {
+  for_each = toset(var.management_subnet_ids)
+  id       = each.value
 }

lambda.tf

@@ -14,7 +14,7 @@ resource "aws_lambda_function" "auto_scaling_lambda" {
 
   environment {
     variables = {
-      TARGET_SUBNET            = var.management_subnet_id
+      TARGET_SUBNETS           = jsonencode({ for subnet in data.aws_subnet.management_subnets : subnet.availability_zone => subnet.id })
       TARGET_SECURITY_GROUP_ID = aws_security_group.management.id
     }
   }

launch_template.tf

@@ -15,7 +15,7 @@ resource "aws_launch_template" "sensor_launch_template" {
   }
 
   network_interfaces {
-    subnet_id             = var.monitoring_subnet_id
+    device_index          = 0
     security_groups       = [aws_security_group.monitoring.id]
     delete_on_termination = true
   }

load_balancer.tf

@@ -1,7 +1,7 @@
 resource "aws_lb" "sensor_lb" {
   name                             = var.sensor_asg_load_balancer_name
   load_balancer_type               = "gateway"
-  subnets                          = [data.aws_subnet.monitoring_subnet.id]
+  subnets                          = [for subnet in data.aws_subnet.monitoring_subnets : subnet.id]
   enable_cross_zone_load_balancing = true
 }
 

modules/iam/lambda/main.tf

@@ -71,11 +71,13 @@ data "aws_iam_policy_document" "lambda_nic_manager_policy" {
       "ec2:CreateNetworkInterface",
       "ec2:CreateTags",
     ]
-    resources = [
-      var.subnet_arn,
-      var.security_group_arn,
-      "arn:aws:ec2:*:*:network-interface/*"
-    ]
+    resources = concat(
+      var.subnet_arns,
+      [
+        var.security_group_arn,
+        "arn:aws:ec2:*:*:network-interface/*"
+      ]
+    )
   }
 }
 

modules/iam/lambda/variables.tf

@@ -8,9 +8,9 @@ variable "sensor_autoscaling_group_arn" {
   type        = string
 }
 
-variable "subnet_arn" {
-  description = "ARN of the subnet where new ENIs should be created (management)"
-  type        = string
+variable "subnet_arns" {
+  description = "ARNs of the subnets where new ENIs should be created (management), one per availability zone"
+  type        = list(string)
 }
 
 variable "security_group_arn" {

scripts/corelight_sensor_asg_nic_manager.py

@@ -1,4 +1,5 @@
 import os
+import json
 from enum import Enum
 
 import boto3
@@ -9,7 +10,7 @@
 
 @dataclass
 class EnvironmentConfig:
-    subnet_id: str
+    subnet_map: dict  # Maps AZ to subnet ID
     security_group_id: str
 
 
@@ -38,7 +39,7 @@ class LifecycleActionResult(Enum):
 
 
 class EnvironmentVariables(Enum):
-    TARGET_SUBNET = "TARGET_SUBNET"
+    TARGET_SUBNETS = "TARGET_SUBNETS"
     TARGET_SECURITY_GROUP_ID = "TARGET_SECURITY_GROUP_ID"
 
 
@@ -135,7 +136,18 @@ def __init__(self, config: EnvironmentConfig, aws_client: AwsClient):
         self.instance_data = {}
 
     def process_event(self, event: Ec2LifecycleHookEvent):
-        network_interface_id = self.aws_client.create_interface(self.config.subnet_id, self.config.security_group_id)
+        # Get the AZ of the instance
+        instance_az = self.instance_data['Placement']['AvailabilityZone']
+        logging.info(f"Instance {event.instance_id} is in AZ {instance_az}")
+
+        # Find the matching management subnet for this AZ
+        if instance_az not in self.config.subnet_map:
+            raise Exception(f"No management subnet configured for AZ {instance_az}. Available AZs: {list(self.config.subnet_map.keys())}")
+
+        target_subnet_id = self.config.subnet_map[instance_az]
+        logging.info(f"Using management subnet {target_subnet_id} for AZ {instance_az}")
+
+        network_interface_id = self.aws_client.create_interface(target_subnet_id, self.config.security_group_id)
         try:
             attachment_resp = self.aws_client.attach_interface(network_interface_id, event.instance_id)
             self.aws_client.modify_attachment_to_delete_on_termination(attachment_resp["AttachmentId"], network_interface_id)
@@ -179,21 +191,29 @@ def lambda_handler(event, context):
     parsed_event: Ec2LifecycleHookEvent = from_aws_event_bridge_json(event)
 
     try:
+        if not lifecycle_event_svc.should_process_event(parsed_event):
+            logging.error("Event validation failed, abandoning lifecycle action")
+            lifecycle_event_svc.complete_lifecycle_action(parsed_event, LifecycleActionResult.ABANDON)
+            return
+
         lifecycle_event_svc.process_event(parsed_event)
         lifecycle_event_svc.complete_lifecycle_action(parsed_event, LifecycleActionResult.CONTINUE)
         logging.info("Lifecycle action completed successfully")
     except Exception as e:
-        logging.info(f"failed to process event: {e}")
-        lifecycle_event_svc.complete_lifecycle_action(parsed_event, LifecycleActionResult.ABANDON)
+        logging.error(f"failed to process event: {e}")
+        try:
+            lifecycle_event_svc.complete_lifecycle_action(parsed_event, LifecycleActionResult.ABANDON)
+        except Exception as complete_error:
+            logging.error(f"Failed to complete lifecycle action with ABANDON: {complete_error}")
         raise e
 
 
 def parse_environment() -> EnvironmentConfig:
-    subnet = os.getenv(EnvironmentVariables.TARGET_SUBNET.value, "")
+    subnets_json = os.getenv(EnvironmentVariables.TARGET_SUBNETS.value, "")
     security_group_id = os.getenv(EnvironmentVariables.TARGET_SECURITY_GROUP_ID.value, "")
 
-    if subnet == "":
-        msg = f"environment variable ${EnvironmentVariables.TARGET_SUBNET.value} is not defined"
+    if subnets_json == "":
+        msg = f"environment variable ${EnvironmentVariables.TARGET_SUBNETS.value} is not defined"
         logging.error(msg)
         raise Exception(msg)
 
@@ -202,4 +222,11 @@ def parse_environment() -> EnvironmentConfig:
         logging.error(msg)
         raise Exception(msg)
 
-    return EnvironmentConfig(subnet_id=subnet, security_group_id=security_group_id)
+    try:
+        subnet_map = json.loads(subnets_json)
+    except json.JSONDecodeError as e:
+        msg = f"Failed to parse TARGET_SUBNETS as JSON: {e}"
+        logging.error(msg)
+        raise Exception(msg)
+
+    return EnvironmentConfig(subnet_map=subnet_map, security_group_id=security_group_id)

scripts/tests/test_lifecycle_event_service.py

@@ -17,7 +17,7 @@
     lifecycle_action_token="87654321-4321-4321-4321-210987654321"
 )
 
-cfg = EnvironmentConfig("foo", "bar")
+cfg = EnvironmentConfig({"us-east-1a": "subnet-foo", "us-east-1b": "subnet-bar"}, "sg-12345")
 aws_client = AwsClient("foo", "bar")
 
 
@@ -90,6 +90,7 @@ def test_process_event_should_raise_exception_on_nic_creation_client_error(mocke
     )
 
     svc = LifecycleEventService(cfg, aws_client)
+    svc.instance_data = {"Placement": {"AvailabilityZone": "us-east-1a"}}
 
     with pytest.raises(botocore.exceptions.ClientError):
         svc.process_event(event)
@@ -108,8 +109,11 @@ def test_process_event_should_raise_exception_and_delete_nic_if_attachment_fails
 
     delete_nic_mocker = mocker.patch.object(aws_client, "delete_interface", return_value=None)
 
+    svc = LifecycleEventService(cfg, aws_client)
+    svc.instance_data = {"Placement": {"AvailabilityZone": "us-east-1a"}}
+
     with pytest.raises(botocore.exceptions.ClientError):
-        LifecycleEventService(cfg, aws_client).process_event(event)
+        svc.process_event(event)
     assert create_nic_mocker.call_count == 1 and \
            attach_interface_mocker.call_count == 1 and \
            delete_nic_mocker.call_count == 1
@@ -133,8 +137,11 @@ def test_process_event_should_raise_exception_and_delete_nic_if_attachment_modif
 
     delete_nic_mocker = mocker.patch.object(aws_client, "delete_interface", return_value=None)
 
+    svc = LifecycleEventService(cfg, aws_client)
+    svc.instance_data = {"Placement": {"AvailabilityZone": "us-east-1a"}}
+
     with pytest.raises(botocore.exceptions.ClientError):
-        LifecycleEventService(cfg, aws_client).process_event(event)
+        svc.process_event(event)
 
     assert create_nic_mocker.call_count == 1 and \
         attach_mocker.call_count == 1 and \