add support for buildkite token from ssm

backend/ssm.go

@@ -0,0 +1,34 @@
+package backend
+
+import (
+	"fmt"
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/ssm"
+	"os"
+	"sync"
+)
+
+var ssmClient *ssm.SSM
+var once sync.Once
+
+func GetSsmClient() *ssm.SSM {
+	once.Do(func() {
+		ssmClient = ssm.New(session.Must(session.NewSession()))
+	})
+	return ssmClient
+}
+
+func RetrieveFromParameterStore(ssmClient *ssm.SSM, key string) string {
+	output, err := ssmClient.GetParameter(&ssm.GetParameterInput{
+		Name:           &key,
+		WithDecryption: aws.Bool(true),
+	})
+
+	if err != nil {
+		fmt.Printf("Error retrieving SSM (%s): %v", key, err)
+		os.Exit(1)
+	}
+
+	return *output.Parameter.Value
+}

lambda/main.go

@@ -29,18 +29,25 @@ func main() {
 
 func Handler(ctx context.Context, evt json.RawMessage) (string, error) {
 	token := os.Getenv("BUILDKITE_AGENT_TOKEN")
+	useSsmString := os.Getenv("BUILDKITE_TOKEN_IN_SSM")
 	backendOpt := os.Getenv("BUILDKITE_BACKEND")
 	queue := os.Getenv("BUILDKITE_QUEUE")
 	clwDimensions := os.Getenv("BUILDKITE_CLOUDWATCH_DIMENSIONS")
 	quietString := os.Getenv("BUILDKITE_QUIET")
 	quiet := quietString == "1" || quietString == "true"
+	useSsm := useSsmString == "true"
 
 	if quiet {
 		log.SetOutput(ioutil.Discard)
 	}
 
 	t := time.Now()
 
+	if useSsm {
+		ssmClient := backend.GetSsmClient()
+		token = backend.RetrieveFromParameterStore(ssmClient, "buildkite_agent_token")
+	}
+
 	userAgent := fmt.Sprintf("buildkite-agent-metrics/%s buildkite-agent-metrics-lambda", version.Version)
 
 	c := collector.Collector{