From fba0ff0fafff7457f64fc74c19ab45995c76e6a4 Mon Sep 17 00:00:00 2001
From: Ryan <ryan@scalesec.com>
Date: Thu, 20 Dec 2018 13:07:06 -0800
Subject: [PATCH] Fixed EC2 Credentials

---
 examples/terraform-ec2/data.tf       |  6 ++++++
 examples/terraform-ec2/iam.tf        |  5 +++++
 examples/terraform-ec2/user-data.tpl |  4 ++--
 main.go                              | 11 +++++++----
 4 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/examples/terraform-ec2/data.tf b/examples/terraform-ec2/data.tf
index d2d1d52..9c61a90 100644
--- a/examples/terraform-ec2/data.tf
+++ b/examples/terraform-ec2/data.tf
@@ -1,3 +1,5 @@
+data "aws_caller_identity" "current" {}
+
 data "aws_ami" "amznlinux" {
   most_recent = true
 
@@ -16,3 +18,7 @@ data "template_file" "user_data" {
     s3_uri  = "${local.s3_uri}"
   }
 }
+
+data "aws_iam_policy" "security_audit" {
+  arn = "arn:aws:iam::${data.aws_caller_identity.current.account_id}:policy/SecurityAudit"
+}
diff --git a/examples/terraform-ec2/iam.tf b/examples/terraform-ec2/iam.tf
index b9e471f..66236f2 100644
--- a/examples/terraform-ec2/iam.tf
+++ b/examples/terraform-ec2/iam.tf
@@ -33,6 +33,11 @@ resource "aws_iam_role" "fedrampup" {
 EOF
 }
 
+resource "aws_iam_role_policy_attachment" "fedrampup_audit_attachment" {
+  role = "${aws_iam_role.fedrampup.name}"
+  policy_arn = "${data.aws_iam_policy.security_audit.arn}"
+}
+
 resource "aws_iam_role_policy_attachment" "fedrampup_attachment" {
   role       = "${aws_iam_role.fedrampup.name}"
   policy_arn = "${aws_iam_policy.fedrampup.arn}"
diff --git a/examples/terraform-ec2/user-data.tpl b/examples/terraform-ec2/user-data.tpl
index 84fa83e..9517b71 100644
--- a/examples/terraform-ec2/user-data.tpl
+++ b/examples/terraform-ec2/user-data.tpl
@@ -18,8 +18,8 @@ mkdir -p /opt/go/src /opt/go/pkg /opt/go/bin
 WRAPPER=/opt/fedrampup-wrapper
 cat << EOF > $WRAPPER
 #!/bin/bash
-AWS_REGION=${aws_region}
-OUTPUT_FILE=${s3_uri}
+export AWS_REGION=${aws_region}
+export OUTPUT_FILE=${s3_uri}
 
 /opt/go/bin/fedrampup
 EOF
diff --git a/main.go b/main.go
index 8ac8c9f..f115830 100644
--- a/main.go
+++ b/main.go
@@ -6,6 +6,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
 	"github.com/aws/aws-sdk-go/aws/defaults"
+	"github.com/aws/aws-sdk-go/aws/ec2metadata"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/s3/s3manager"
 	"io/ioutil"
@@ -53,11 +54,13 @@ func main() {
 func GetSession() *session.Session {
 	var creds *credentials.Credentials
 	sess := session.Must(session.NewSession())
-
-	if len(os.Getenv("AWS_ACCESS_KEY_ID")) > 0 {
-		creds = credentials.NewEnvCredentials()
+	meta := ec2metadata.New(sess)
+	if meta.Available() {
+		creds = credentials.NewCredentials(&ec2rolecreds.EC2RoleProvider{
+			Client: meta,
+		})
 	} else {
-		creds = credentials.NewCredentials(&ec2rolecreds.EC2RoleProvider{})
+		creds = credentials.NewEnvCredentials()
 	}
 	if _, err := creds.Get(); err != nil {
 		log.Fatal(err)