Add possibility to use public IP

README.md

@@ -36,7 +36,7 @@ Once you have your ASG set up, you can just invoke this module and point to it:
 module "clever_name_autoscale_dns" {
   source  = "meltwater/asg-dns-handler/aws"
   version = "x.y.z"
-
+  # use_public_ip = true
   autoscale_handler_unique_identifier = "clever_name"
   autoscale_route53zone_arn           = "ABCDEFGHIJ123"
   vpc_name                            = "my_vpc"

lambda/autoscale/autoscale.py

@@ -2,6 +2,7 @@
 import logging
 import boto3
 import sys
+import os
 
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
@@ -15,20 +16,22 @@
 LIFECYCLE_KEY = "LifecycleHookName"
 ASG_KEY = "AutoScalingGroupName"
 
-# Fetches private IP of an instance via EC2 API
-def fetch_private_ip_from_ec2(instance_id):
-    logger.info("Fetching private IP for instance-id: %s", instance_id)
-
+# Fetches IP of an instance via EC2 API
+def fetch_ip_from_ec2(instance_id):
+    logger.info("Fetching IP for instance-id: %s", instance_id)
     ec2_response = ec2.describe_instances(InstanceIds=[instance_id])
-    ip_address = ec2_response['Reservations'][0]['Instances'][0]['NetworkInterfaces'][0]['PrivateIpAddress']
-
-    logger.info("Found private IP for instance-id %s: %s", instance_id, ip_address)
+    if 'use_public_ip' in os.environ and os.environ['use_public_ip'] == "true":
+        ip_address = ec2_response['Reservations'][0]['Instances'][0]['PublicIpAddress']
+        logger.info("Found public IP for instance-id %s: %s", instance_id, ip_address)
+    else:
+        ip_address = ec2_response['Reservations'][0]['Instances'][0]['PrivateIpAddress']
+        logger.info("Found private IP for instance-id %s: %s", instance_id, ip_address)
 
     return ip_address
 
-# Fetches private IP of an instance via route53 API
-def fetch_private_ip_from_route53(hostname, zone_id):
-    logger.info("Fetching private IP for hostname: %s", hostname)
+# Fetches IP of an instance via route53 API
+def fetch_ip_from_route53(hostname, zone_id):
+    logger.info("Fetching IP for hostname: %s", hostname)
 
     ip_address = route53.list_resource_record_sets(
         HostedZoneId=zone_id,
@@ -37,7 +40,7 @@ def fetch_private_ip_from_route53(hostname, zone_id):
         MaxItems='1'
     )['ResourceRecordSets'][0]['ResourceRecords'][0]['Value']
 
-    logger.info("Found private IP for hostname %s: %s", hostname, ip_address)
+    logger.info("Found IP for hostname %s: %s", hostname, ip_address)
 
     return ip_address
 
@@ -99,8 +102,11 @@ def update_record(zone_id, ip, hostname, operation):
     )
 
 # Processes a scaling event
-# Builds a hostname from tag metadata, fetches a private IP, and updates records accordingly
+# Builds a hostname from tag metadata, fetches a IP, and updates records accordingly
 def process_message(message):
+    if 'LifecycleTransition' not in message:
+        logger.info("Processing %s event", message['Event'])
+        return
     logger.info("Processing %s event", message['LifecycleTransition'])
 
     if message['LifecycleTransition'] == "autoscaling:EC2_INSTANCE_LAUNCHING":
@@ -117,13 +123,13 @@ def process_message(message):
     hostname = build_hostname(hostname_pattern, instance_id)
 
     if operation == "UPSERT":
-        private_ip = fetch_private_ip_from_ec2(instance_id)
+        ip = fetch_ip_from_ec2(instance_id)
 
         update_name_tag(instance_id, hostname)
     else:
-        private_ip = fetch_private_ip_from_route53(hostname, zone_id)
+        ip = fetch_ip_from_route53(hostname, zone_id)
 
-    update_record(zone_id, private_ip, hostname, operation)
+    update_record(zone_id, ip, hostname, operation)
 
 # Picks out the message from a SNS message and deserializes it
 def process_record(record):
@@ -147,14 +153,15 @@ def lambda_handler(event, context):
             InstanceId = message['EC2InstanceId'],
             LifecycleActionToken = message['LifecycleActionToken'],
             LifecycleActionResult = 'CONTINUE'
-        
+
         )
-        logger.info("ASG action complete: %s", response)    
+        logger.info("ASG action complete: %s", response)
     else :
-        logger.error("No valid JSON message")        
+        logger.error("No valid JSON message")
 
 # if invoked manually, assume someone pipes in a event json
 if __name__ == "__main__":
     logging.basicConfig()
 
     lambda_handler(json.load(sys.stdin), None)
+

main.tf

@@ -113,6 +113,11 @@ resource "aws_lambda_function" "autoscale_handling" {
   runtime          = "python2.7"
   source_code_hash = filebase64sha256(data.archive_file.autoscale.output_path)
   description      = "Handles DNS for autoscaling groups by receiving autoscaling notifications and setting/deleting records from route53"
+  environment {
+    variables = {
+      "use_public_ip" = var.use_public_ip
+    }
+  }
 }
 
 resource "aws_lambda_permission" "autoscale_handling" {

variables.tf

@@ -6,6 +6,11 @@ variable "vpc_name" {
   description = "The name of the VPC"
 }
 
+variable "use_public_ip" {
+  description = "Use public IP instead of private"
+  default     = false
+}
+
 variable "autoscale_route53zone_arn" {
   description = "The ARN of route53 zone associated with autoscaling group"
 }