1 files changed, 75 insertions, 9 deletions

diff --git a/app/controllers/api_credentials.go b/app/controllers/api_credentials.go
index cd1a912..d53565e 100644
--- a/app/controllers/api_credentials.go
+++ b/app/controllers/api_credentials.go
@@ -1,7 +1,9 @@
 package controllers
 
 import (
+        "bytes"
         "net/http"
+        "text/template"
         "time"
 
         "code.crute.us/mcrute/cloud-identity-broker/cloud/aws"
@@ -23,8 +25,33 @@ type jsonCredential struct {
         SecretAccessKey *string    `json:"secret_key"`
         SessionToken    *string    `json:"session_token"`
         Expiration      *time.Time `json:"expiration"`
+        ShortName       string     `json:"-"`
 }
 
+var (
+        bashTemplate = template.Must(template.New("").Parse(
+                `export AWS_CREDS_EXPIRATION="{{ .Expiration }}"
+export AWS_ACCESS_KEY_ID="{{ .AccessKeyId }}"
+export AWS_SECRET_ACCESS_KEY="{{ .SecretAccessKey }}"
+export AWS_SESSION_TOKEN="{{ .SessionToken }}"
+`))
+
+        pslTemplate = template.Must(template.New("").Parse(
+                `Set-Item -path env:AWS_CREDS_EXPIRATION -value '{{ .Expiration }}'
+Set-Item -path env:AWS_ACCESS_KEY_ID -value '{{ .AccessKeyId }}'
+Set-Item -path env:AWS_SECRET_ACCESS_KEY -value '{{ .SecretAccessKey }}'
+Set-Item -path env:AWS_SESSION_TOKEN -value '{{ .SessionToken }}'
+`))
+
+        iniTemplate = template.Must(template.New("").Parse(
+                `[profile {{ .ShortName }}]
+aws_access_key_id={{ .AccessKeyId }}
+aws_secret_access_key={{ .SecretAccessKey }}
+aws_session_token={{ .SessionToken }}
+expiration={{ .Expiration }}
+`))
+)
+
 type APICredentialsHandler struct {
         *AWSAPI
 }
@@ -32,29 +59,32 @@ type APICredentialsHandler struct {
 func NewAPICredentialsHandler(a *AWSAPI) echo.HandlerFunc {
         al := &APICredentialsHandler{a}
         h := &controller.ContentTypeNegotiatingHandler{
-                DefaultHandler: al.Handle,
+                DefaultHandler: al.HandleJSONV1,
                 Handlers: map[string]echo.HandlerFunc{
-                        contentTypeV1: al.Handle,
-                        contentTypeV2: al.Handle,
+                        contentTypeV1:              al.HandleJSONV1,
+                        contentTypeV2:              al.HandleJSONV1,
+                        contentTypeV2AWSBash:       al.HandleBashV2,
+                        contentTypeV2AWSPowershell: al.HandlePSLV2,
+                        contentTypeV2AWSConfig:     al.HandleINIV2,
                 },
         }
         return h.Handle
 }
 
-func (h *APICredentialsHandler) Handle(c echo.Context) error {
+func (h *APICredentialsHandler) getAWSCredential(c echo.Context) (*jsonCredential, error) {
         rc, err := h.GetContext(c) // Does authorization checks
         if err != nil {
-                return err
+                return nil, err
         }
 
         region := c.Param("region")
         creds, err := rc.AWS.AssumeRole(rc.Principal.Username, &region)
         if err != nil {
                 if aws.IsRegionNotExist(err) {
-                        return echo.NotFoundHandler(c)
+                        return nil, echo.NotFoundHandler(c)
                 }
                 c.Logger().Errorf("Error retrieving credentials: %w", err)
-                return echo.ErrInternalServerError
+                return nil, echo.ErrInternalServerError
         }
 
         c.Logger().Infof(
@@ -68,10 +98,46 @@ func (h *APICredentialsHandler) Handle(c echo.Context) error {
 
         c.Response().Header().Set("Expires", creds.Expiration.Add(-5*time.Minute).Format(time.RFC1123))
 
-        return c.JSON(http.StatusOK, &jsonCredential{
+        return &jsonCredential{
                 AccessKeyId:     creds.AccessKeyId,
                 SecretAccessKey: creds.SecretAccessKey,
                 SessionToken:    creds.SessionToken,
                 Expiration:      creds.Expiration,
-        })
+                ShortName:       rc.Account.ShortName,
+        }, nil
+}
+
+func (h *APICredentialsHandler) renderTemplate(c echo.Context, t *template.Template, ct string) error {
+        creds, err := h.getAWSCredential(c)
+        if err != nil {
+                return err
+        }
+
+        buf := &bytes.Buffer{}
+        if err = t.Execute(buf, creds); err != nil {
+                return echo.ErrInternalServerError
+        }
+
+        return c.Blob(http.StatusOK, ct, buf.Bytes())
+}
+
+func (h *APICredentialsHandler) HandleJSONV1(c echo.Context) error {
+        creds, err := h.getAWSCredential(c)
+        if err != nil {
+                return err
+        }
+
+        return c.JSON(http.StatusOK, creds)
+}
+
+func (h *APICredentialsHandler) HandleBashV2(c echo.Context) error {
+        return h.renderTemplate(c, bashTemplate, contentTypeV2AWSBash)
+}
+
+func (h *APICredentialsHandler) HandlePSLV2(c echo.Context) error {
+        return h.renderTemplate(c, pslTemplate, contentTypeV2AWSPowershell)
+}
+
+func (h *APICredentialsHandler) HandleINIV2(c echo.Context) error {
+        return h.renderTemplate(c, iniTemplate, contentTypeV2AWSConfig)
 }