Merge pull request #7 from ApplauseOSS/banka/support-assume-role

Add assume-role support

Author: kbanka

README.md

@@ -25,3 +25,9 @@ It can also take an optional flag to control the number of parallel workers:
 ```bash
 $ decrypt-and-start -p 20 -- some other program
 ```
+
+Tool can also assume other role for kms access
+
+```bash
+$ decrypt-and-start --assume-role arn:aws:iam::XXXXXXXXX:role/YYYY some other program
+```

VERSION

@@ -1 +1 @@
-0.3.0
+0.4.0

decrypt-and-start.go

@@ -28,10 +28,12 @@ func Exec() {
 
 func main() {
 	var workerCount int
+	var assumedRole string
 	flag.IntVar(&workerCount, "p", 10, "number of parallel workers (defaults to 10)")
+	flag.StringVar(&assumedRole, "assume-role", "", "Arn of role to assume for variables decryption")
 	flag.Parse()
 	workerPool := lib.NewWorkerPool(workerCount)
-	workerPool.Start()
+	workerPool.Start(assumedRole)
 	// Put encrypted env vars in queue for workers to process
 	go func() {
 		for _, e := range os.Environ() {

lib/aws_encryption_sdk/kms_helper.go

@@ -7,6 +7,7 @@ import (
 	"encoding/binary"
 	"errors"
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/kms"
 	"golang.org/x/crypto/hkdf"
@@ -17,12 +18,17 @@ type KmsHelper struct {
 	client *kms.KMS
 }
 
-func NewKmsHelper(region string) *KmsHelper {
+func NewKmsHelper(region string, assumedRole string) *KmsHelper {
 	k := &KmsHelper{}
 	// Set up AWS KMS session
 	conf := aws.NewConfig().WithRegion(region)
 	sess := session.Must(session.NewSession(conf))
-	k.client = kms.New(sess)
+	if assumedRole != "" {
+		creds := stscreds.NewCredentials(sess, assumedRole)
+		k.client = kms.New(sess, &aws.Config{Credentials: creds})
+	} else {
+		k.client = kms.New(sess)
+	}
 	return k
 }
 

lib/worker.go

@@ -27,10 +27,10 @@ func NewWorkerPool(count int) *WorkerPool {
 	return w
 }
 
-func (w *WorkerPool) Start() {
+func (w *WorkerPool) Start(assumedRole string) {
 	for i := 0; i < w.workerCount; i++ {
 		go func() {
-			kmsHelper := enc_sdk.NewKmsHelper(GetRegion())
+			kmsHelper := enc_sdk.NewKmsHelper(GetRegion(), assumedRole)
 			for {
 				env, ok := <-w.InChan
 				if env != nil {