From 90a46bd02cd3dd6176f89bf205184e043c083376 Mon Sep 17 00:00:00 2001
From: Jeff Storey <jeffrey.storey@capitalone.com>
Date: Wed, 23 Mar 2016 08:43:20 -0400
Subject: [PATCH] allow assume role before running terraform commands

---
 DOCS.md | 20 ++++++++++++++++++++
 main.go | 34 ++++++++++++++++++++++++++++++++--
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/DOCS.md b/DOCS.md
index 6bd526e..92c1e16 100644
--- a/DOCS.md
+++ b/DOCS.md
@@ -8,6 +8,7 @@ Use the Terraform plugin to apply the infrastructure configuration contained wit
  <key>=<value>` option.
 * `ca_cert` - ca cert to add to your environment to allow terraform to use internal/private resources
 * `sensitive` (default: `false`) - Whether or not to suppress terraform commands to stdout.
+* `role_arn_to_assume` - A role to assume before running the terraform commands
 
 The following is a sample Terraform configuration in your .drone.yml file:
 
@@ -73,3 +74,22 @@ deploy:
         app_name: my-project
         app_version: 1.0.0
 ```
+
+## Assume Role ARN
+You may want to assume another role before running the terraform commands. This is useful for cross account access, where a central account ahs privileges to assume roles in other accounts. Using the current credentials, this role will be assumed and exported to environment variables.  See [the discussion](https://github.com/hashicorp/terraform/issues/1275) in the Terraform issues.
+
+```yaml
+deploy:
+  terraform:
+    plan: false
+    remote:
+      backend: S3
+      config:
+        bucket: my-terraform-config-bucket
+        key: tf-states/my-project
+        region: us-east-1
+    vars:
+      app_name: my-project
+      app_version: 1.0.0
+    role_arn_to_assume: arn:aws:iam::account-of-role-to-assume:role/name-of-role
+```
diff --git a/main.go b/main.go
index 4e4d0ab..cc13f49 100644
--- a/main.go
+++ b/main.go
@@ -6,8 +6,12 @@ import (
 	"os"
 	"os/exec"
 	"strings"
-
-	"github.com/drone/drone-plugin-go/plugin"
+    "time"
+    "github.com/aws/aws-sdk-go/aws/credentials"
+    "github.com/aws/aws-sdk-go/aws/credentials/stscreds"
+    "github.com/aws/aws-sdk-go/aws/session"
+    "github.com/aws/aws-sdk-go/service/sts"
+    "github.com/drone/drone-plugin-go/plugin"
 )
 
 var (
@@ -20,6 +24,7 @@ type terraform struct {
 	Vars      map[string]string `json:"vars"`
 	Cacert    string            `json:"ca_cert"`
 	Sensitive bool              `json:"sensitive"`
+	RoleARN	  string            `json:"role_arn_to_assume"`
 }
 
 type remote struct {
@@ -37,6 +42,10 @@ func main() {
 	plugin.Param("vargs", &vargs)
 	plugin.MustParse()
 
+	if vargs.RoleARN != "" {
+		assumeRole(vargs.RoleARN)
+	}
+
 	var commands []*exec.Cmd
 	remote := vargs.Remote
 	if vargs.Cacert != "" {
@@ -125,6 +134,27 @@ func applyCommand() *exec.Cmd {
 	)
 }
 
+func assumeRole(roleArn string) {
+    client := sts.New(session.New())
+    duration := time.Hour * 1
+    stsProvider := &stscreds.AssumeRoleProvider{
+        Client: client,
+        Duration: duration,
+        RoleARN: roleArn,
+        RoleSessionName: "drone",
+    }
+
+    value, err := credentials.NewCredentials(stsProvider).Get()
+	if err != nil {
+		fmt.Println("Error assuming role!")
+		fmt.Println(err)
+		os.Exit(1)
+	}
+    os.Setenv("AWS_ACCESS_KEY_ID",value.AccessKeyID)
+    os.Setenv("AWS_SECRET_ACCESS_KEY",value.SecretAccessKey)
+    os.Setenv("AWS_SESSION_TOKEN",value.SessionToken)
+}
+
 func trace(cmd *exec.Cmd) {
 	fmt.Println("$", strings.Join(cmd.Args, " "))
 }