github.com/hazelops/ize@v1.1.12-0.20230915191306-97d7c0e48f11/internal/aws/utils/utils.go (about)

     1  package utils
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/aws/aws-sdk-go/aws/awserr"
     6  	"github.com/aws/aws-sdk-go/aws/credentials"
     7  	"github.com/sirupsen/logrus"
     8  	"os"
     9  	"time"
    10  
    11  	"github.com/aws/aws-sdk-go/aws"
    12  	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
    13  	"github.com/aws/aws-sdk-go/aws/session"
    14  	"github.com/aws/aws-sdk-go/service/iam"
    15  	"github.com/aws/aws-sdk-go/service/sts"
    16  	"gopkg.in/ini.v1"
    17  )
    18  
    19  const (
    20  	path = "/.aws/credentials-mfa"
    21  )
    22  
    23  type SessionConfig struct {
    24  	Region  string
    25  	Profile string
    26  }
    27  
    28  func GetSession(c *SessionConfig) (*session.Session, error) {
    29  	upd := false
    30  
    31  	config := aws.NewConfig().WithRegion(c.Region).WithCredentials(credentials.NewSharedCredentials("", c.Profile))
    32  	sess, err := session.NewSessionWithOptions(session.Options{
    33  		Config: *config,
    34  	})
    35  	if err != nil {
    36  		return nil, err
    37  	}
    38  
    39  	devices, err := iam.New(sess).ListMFADevices(&iam.ListMFADevicesInput{})
    40  	if aerr, ok := err.(awserr.Error); ok {
    41  		switch aerr.Code() {
    42  		case "SharedCredsLoad":
    43  			logrus.Error(err)
    44  			return nil, fmt.Errorf("AWS_PROFILE is not set. Please set it via AWS_PROFILE env var,--aws-profile flag or aws_profile config entry in ize.toml")
    45  		default:
    46  			return nil, err
    47  		}
    48  	}
    49  
    50  	if len(devices.MFADevices) == 0 {
    51  		return sess, nil
    52  	}
    53  
    54  	home, _ := os.UserHomeDir()
    55  	filePath := home + path
    56  
    57  	credFile, err := ini.Load(filePath)
    58  	if err != nil {
    59  		credFile = ini.Empty(ini.LoadOptions{})
    60  		upd = true
    61  	}
    62  
    63  	var sect *ini.Section
    64  	var exp *ini.Key
    65  
    66  	if !upd {
    67  		sect, err = credFile.GetSection(fmt.Sprintf("%s-mfa", c.Profile))
    68  		if err != nil {
    69  			upd = true
    70  		}
    71  	}
    72  
    73  	if !upd {
    74  		if len(sect.KeyStrings()) != 4 {
    75  			upd = true
    76  		}
    77  	}
    78  
    79  	if !upd {
    80  		exp, err = sect.GetKey("token_expiration")
    81  		if err != nil {
    82  			upd = true
    83  		}
    84  	}
    85  
    86  	if !upd {
    87  		timeExp, err := time.Parse("2006-01-02T15:04:05Z07:00", exp.String())
    88  		if err != nil {
    89  			upd = true
    90  		}
    91  
    92  		if timeExp.Before(time.Now().UTC()) {
    93  			upd = true
    94  		}
    95  	}
    96  
    97  	if upd {
    98  		cred, err := getNewToken(sess, devices.MFADevices[0].SerialNumber)
    99  		if err != nil {
   100  			return nil, err
   101  		}
   102  
   103  		err = writeCredsToFile(cred, credFile, filePath, c.Profile)
   104  		if err != nil {
   105  			return nil, err
   106  		}
   107  	}
   108  
   109  	sess, err = session.NewSessionWithOptions(
   110  		session.Options{
   111  			Config:            *aws.NewConfig().WithRegion(c.Region),
   112  			Profile:           fmt.Sprintf("%s-mfa", c.Profile),
   113  			SharedConfigFiles: []string{filePath},
   114  		},
   115  	)
   116  	if err != nil {
   117  		return nil, err
   118  	}
   119  
   120  	return sess, nil
   121  }
   122  
   123  func GetTestSession(c *SessionConfig) (*session.Session, error) {
   124  	sess, err := session.NewSessionWithOptions(
   125  		session.Options{
   126  			Config:  *aws.NewConfig().WithRegion(c.Region),
   127  			Profile: c.Profile,
   128  		},
   129  	)
   130  	if err != nil {
   131  		return nil, err
   132  	}
   133  
   134  	return sess, nil
   135  }
   136  
   137  func getNewToken(sess *session.Session, sn *string) (*sts.Credentials, error) {
   138  	stsSvc := sts.New(sess)
   139  
   140  	token, err := stscreds.StdinTokenProvider()
   141  	if err != nil {
   142  		return nil, err
   143  	}
   144  
   145  	out, err := stsSvc.GetSessionToken(&sts.GetSessionTokenInput{
   146  		SerialNumber: sn,
   147  		TokenCode:    &token,
   148  	})
   149  
   150  	if err != nil {
   151  		return nil, err
   152  	}
   153  
   154  	return out.Credentials, nil
   155  }
   156  
   157  func writeCredsToFile(creds *sts.Credentials, f *ini.File, filepath, profile string) error {
   158  	sect, err := f.NewSection(fmt.Sprintf("%s-mfa", profile))
   159  	if err != nil {
   160  		return err
   161  	}
   162  
   163  	_, err = sect.NewKey("aws_access_key_id", *creds.AccessKeyId)
   164  	if err != nil {
   165  		return err
   166  	}
   167  	_, err = sect.NewKey("aws_secret_access_key", *creds.SecretAccessKey)
   168  	if err != nil {
   169  		return err
   170  	}
   171  	_, err = sect.NewKey("aws_session_token", *creds.SessionToken)
   172  	if err != nil {
   173  		return err
   174  	}
   175  	_, err = sect.NewKey("token_expiration", creds.Expiration.Format("2006-01-02T15:04:05Z07:00"))
   176  	if err != nil {
   177  		return err
   178  	}
   179  
   180  	err = f.SaveTo(filepath)
   181  	if err != nil {
   182  		return err
   183  	}
   184  
   185  	return nil
   186  }