From 547c45ac398f937e38ecde2c3f535e5a7ad75f0e Mon Sep 17 00:00:00 2001
From: Jose Gaspar <jad.gaspar@gmail.com>
Date: Mon, 29 Jan 2018 10:01:08 +0000
Subject: [PATCH] add environment variable supporting different pem files

Instead of having a encrypted pem file hardcoded, this PR introduces a
environment variable CIPHER_FILE that can be used for different accounts
e.g. dev, staging, prod.
---
 README.md              |  4 +--
 lambda/main.py         |  8 +++---
 terraform/main.tf      | 60 ++++++++++++++++++++++++------------------
 terraform/variables.tf |  5 ++++
 4 files changed, 47 insertions(+), 30 deletions(-)

diff --git a/README.md b/README.md
index 676f28a..61956b4 100644
--- a/README.md
+++ b/README.md
@@ -36,7 +36,7 @@ aws kms encrypt --key-id KEY_FROM_STEP_1 --plaintext file://your_private_key.pem
   }
   ```
 
-4. Copy just the CiphertextBlob value into a new file and store it in the same directory as the Lambda function; this is required so it can be packaged up with the function itself. I’ve used encrypted_pem.txt as the file name in my example, given the encrypted object is a certificate and private key, which is commonly name with the .pem file extension.
+4. Copy just the CiphertextBlob value into a new file and store it in the same directory as the Lambda function; this is required so it can be packaged up with the function itself. I’ve used encrypted_PROD_pem.txt as the file name in my example, given the encrypted object is a certificate and private key, which is commonly name with the .pem file extension. Lambda supports environment variables wich we can use to pass diferent CiphertextBlob depending on the account you are working on e.g. dev, staging, prod.
 
 ***Note*** the CiphertextBlob output is base64 encoded by the AWS CLI unless you send the output to a binary file using:
 
@@ -48,7 +48,7 @@ To test decryption, you can try something like this:
 
 ```
 aws --profile <your profile name here> kms decrypt --ciphertext-blob \n
-fileb://<(cat encrypted_pem.txt | base64 -D) --output text \n
+fileb://<(cat encrypted_PROD_pem.txt | base64 -D) --output text \n
 --query Plaintext | base64 -D
 ```
 ***Note*** - On a Mac, use `-D`, on any other *nix environment, use `-d`.
diff --git a/lambda/main.py b/lambda/main.py
index 811ba87..ccf951a 100644
--- a/lambda/main.py
+++ b/lambda/main.py
@@ -45,10 +45,10 @@ def get_instance_id(event):
         LOGGER.error(err)
         return False
 
-def get_pem():
+def get_pem(cipherfile):
     """Decrypt the Ciphertext Blob to get USERNAME's pem file"""
     try:
-        with open('encrypted_pem.txt', 'r') as encrypted_pem:
+        with open('cipherfile', 'r') as encrypted_pem:
             pem_file = encrypted_pem.read()
 
         kms = boto3.client('kms', region_name=REGION)
@@ -61,9 +61,11 @@ def handle(event, _context):
     """Lambda Handler"""
     log_event(event)
 
+    cipherfile = os.environ['CIPHER_FILE']
+
     # If you're using a self signed certificate change
     # the ssl_verify argument to False
-    with chef.ChefAPI(CHEF_SERVER_URL, get_pem(), USERNAME, ssl_verify=VERIFY_SSL):
+    with chef.ChefAPI(CHEF_SERVER_URL, get_pem(cipherfile), USERNAME, ssl_verify=VERIFY_SSL):
         instance_id = get_instance_id(event)
         try:
             search = chef.Search('node', 'ec2_instance_id:' + instance_id)
diff --git a/terraform/main.tf b/terraform/main.tf
index d7ad92f..ed9cff9 100644
--- a/terraform/main.tf
+++ b/terraform/main.tf
@@ -14,9 +14,10 @@ provider "aws" {
 
 # Lambda Role with Required Policy
 resource "aws_iam_role_policy" "lambda_policy" {
-    name = "chef_node_cleanup_lambda"
-    role = "${aws_iam_role.lambda_role.id}"
-    policy = <<EOF
+  name = "chef_node_cleanup_lambda"
+  role = "${aws_iam_role.lambda_role.id}"
+
+  policy = <<EOF
 {
   "Version": "2012-10-17",
   "Statement": [
@@ -35,8 +36,9 @@ EOF
 }
 
 resource "aws_iam_role" "lambda_role" {
-    name = "chef_node_cleanup_lambda"
-    assume_role_policy = <<EOF
+  name = "chef_node_cleanup_lambda"
+
+  assume_role_policy = <<EOF
 {
   "Version": "2012-10-17",
   "Statement": [
@@ -52,36 +54,44 @@ resource "aws_iam_role" "lambda_role" {
 }
 EOF
 }
+
 output "lambda_role_arn" {
   value = "${aws_iam_role.lambda_role.name}"
 }
 
 # Lambda Function
 resource "aws_lambda_function" "lambda_function" {
-    filename = "lambda_function_payload.zip"
-    function_name = "chef_node_cleanup"
-    role = "${aws_iam_role.lambda_role.arn}"
-    handler = "main.handle"
-    description = "Automatically delete nodes from Chef Server on termination"
-    memory_size = 128
-    runtime = "python2.7"
-    timeout = 5
-    source_code_hash = "${base64encode(sha256(file("lambda_function_payload.zip")))}"
+  filename         = "lambda_function_payload.zip"
+  function_name    = "chef_node_cleanup"
+  role             = "${aws_iam_role.lambda_role.arn}"
+  handler          = "main.handle"
+  description      = "Automatically delete nodes from Chef Server on termination"
+  memory_size      = 128
+  runtime          = "python2.7"
+  timeout          = 5
+  source_code_hash = "${base64encode(sha256(file("lambda_function_payload.zip")))}"
+
+  environment {
+    variables = {
+      CIPHER_FILE = "encrypted_${var.tags["environment"]}_pem.txt"
+    }
+  }
 }
 
 resource "aws_lambda_permission" "allow_cloudwatch" {
-    statement_id = "AllowExecutionFromCloudWatch"
-    action = "lambda:InvokeFunction"
-    function_name = "${aws_lambda_function.lambda_function.arn}"
-    principal = "events.amazonaws.com"
-    source_arn = "${aws_cloudwatch_event_rule.instance_termination.arn}"
+  statement_id  = "AllowExecutionFromCloudWatch"
+  action        = "lambda:InvokeFunction"
+  function_name = "${aws_lambda_function.lambda_function.arn}"
+  principal     = "events.amazonaws.com"
+  source_arn    = "${aws_cloudwatch_event_rule.instance_termination.arn}"
 }
 
 # CloudWatch Event Rule and Event Target
 resource "aws_cloudwatch_event_rule" "instance_termination" {
-  depends_on = ["aws_iam_role.lambda_role"] # we need the Lambda arn to exist
-  name = "Chef_Node_Cleanup_Lambda"
+  depends_on  = ["aws_iam_role.lambda_role"]                                       # we need the Lambda arn to exist
+  name        = "Chef_Node_Cleanup_Lambda"
   description = "Trigger the chef_node_cleanup Lambda when an instance terminates"
+
   event_pattern = <<PATTERN
   {
     "source": [ "aws.ec2" ],
@@ -94,8 +104,8 @@ PATTERN
 }
 
 resource "aws_cloudwatch_event_target" "lambda" {
-  depends_on = ["aws_iam_role.lambda_role"] # we need the Lambda arn to exist
-  rule = "${aws_cloudwatch_event_rule.instance_termination.name}"
-  target_id = "chef_node_cleanup"
-  arn = "${aws_lambda_function.lambda_function.arn}"
+  depends_on = ["aws_iam_role.lambda_role"]                             # we need the Lambda arn to exist
+  rule       = "${aws_cloudwatch_event_rule.instance_termination.name}"
+  target_id  = "chef_node_cleanup"
+  arn        = "${aws_lambda_function.lambda_function.arn}"
 }
diff --git a/terraform/variables.tf b/terraform/variables.tf
index 8e7eab4..5487647 100644
--- a/terraform/variables.tf
+++ b/terraform/variables.tf
@@ -11,3 +11,8 @@
 variable "region" {
   default = "us-west-2"
 }
+
+variable "tags" {
+  type    = "map"
+  default = {}
+}