<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"time"

	"github.com/99designs/aws-vault/keyring"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/sts"
)

const (
	MaxSessionDuration    = time.Hour * 36
	MinSessionDuration    = time.Minute * 15
	MinAssumeRoleDuration = time.Minute * 15
	MaxAssumeRoleDuration = time.Hour

	DefaultSessionDuration    = time.Hour * 4
	DefaultAssumeRoleDuration = time.Minute * 15
)

type stsClient interface {
	AssumeRole(input *sts.AssumeRoleInput) (*sts.AssumeRoleOutput, error)
	GetSessionToken(input *sts.GetSessionTokenInput) (*sts.GetSessionTokenOutput, error)
}

type VaultOptions struct {
	SessionDuration    time.Duration
	AssumeRoleDuration time.Duration
	ExpiryWindow       time.Duration
	WriteEnv           bool
	MfaToken           string
}

func (o VaultOptions) Validate() error {
	if o.SessionDuration < time.Minute*15 {
		return errors.New("Minimum session duration is 15 minutes")
	} else if o.SessionDuration > time.Hour*36 {
		return errors.New("Maximum session duration is 36 hours")
	}
	return nil
}

func (o VaultOptions) ApplyDefaults() VaultOptions {
	if o.AssumeRoleDuration == 0 {
		o.AssumeRoleDuration = DefaultAssumeRoleDuration
	}
	if o.SessionDuration == 0 {
		o.SessionDuration = DefaultSessionDuration
	}
	return o
}

type VaultProvider struct {
	credentials.Expiry
	VaultOptions
	profile  string
	expires  time.Time
	keyring  keyring.Keyring
	profiles profiles
	session  *sts.Credentials
	client   stsClient
}

func NewVaultProvider(k keyring.Keyring, profile string, opts VaultOptions) (*VaultProvider, error) {
	opts = opts.ApplyDefaults()
	if err := opts.Validate(); err != nil {
		return nil, err
	}
	profiles, err := parseProfiles()
	if err != nil {
		return nil, err
	}
	return &VaultProvider{
		VaultOptions: opts,
		keyring:      k,
		profile:      profile,
		profiles:     profiles,
	}, nil
}

func (p *VaultProvider) Retrieve() (credentials.Value, error) {
	creds, err := p.getMasterCreds()
	if err != nil {
		return credentials.Value{}, err
	}

	session, err := p.getCachedSession()
	if err != nil {
		if err == keyring.ErrKeyNotFound {
			log.Println("Session not found in keyring")
		} else {
			log.Println(err)
		}

		session, err = p.getSessionToken(&creds)
		if err != nil {
			return credentials.Value{}, err
		}

		bytes, err := json.Marshal(session)
		if err != nil {
			return credentials.Value{}, err
		}

		log.Printf("Writing session for %s to keyring", p.profile)
		p.keyring.Set(keyring.Item{
			Key:       sessionKey(p.profile),
			Label:     "aws-vault session for " + p.profile,
			Data:      bytes,
			TrustSelf: true,
		})
	}

	log.Printf("Using session, expires in %s", session.Expiration.Sub(time.Now()).String())

	window := p.ExpiryWindow
	if window == 0 {
		window = p.SessionDuration - (p.SessionDuration / 3)
	}

	p.SetExpiration(*session.Expiration, window)
	p.expires = *session.Expiration

	if role, ok := p.profiles[p.profile]["role_arn"]; ok {
		session, err = p.assumeRole(session, role)
		if err != nil {
			return credentials.Value{}, err
		}

		log.Printf("Role token expires in %s", session.Expiration.Sub(time.Now()))
	}

	value := credentials.Value{
		AccessKeyID:     *session.AccessKeyId,
		SecretAccessKey: *session.SecretAccessKey,
		SessionToken:    *session.SessionToken,
	}

	return value, nil
}

func sessionKey(profile string) string {
	return profile + " session"
}

func (p *VaultProvider) getCachedSession() (session sts.Credentials, err error) {
	item, err := p.keyring.Get(sessionKey(p.profile))
	if err != nil {
		return session, err
	}

	if err = json.Unmarshal(item.Data, &session); err != nil {
		return session, err
	}

	if session.Expiration.Before(time.Now()) {
		return session, errors.New("Session is expired")
	}

	return
}

func (p *VaultProvider) getMasterCreds() (credentials.Value, error) {
	source := p.profiles.sourceProfile(p.profile)

	provider := credentials.NewChainCredentials([]credentials.Provider{
		&credentials.EnvProvider{},
		&credentials.SharedCredentialsProvider{Filename: "", Profile: p.profile},
		&KeyringProvider{Keyring: p.keyring, Profile: source},
	})

	return provider.Get()
}

func (p *VaultProvider) getSessionToken(creds *credentials.Value) (sts.Credentials, error) {
	params := &sts.GetSessionTokenInput{
		DurationSeconds: aws.Int64(int64(p.SessionDuration.Seconds())),
	}

	if mfa, ok := p.profiles[p.profile]["mfa_serial"]; ok {
		params.SerialNumber = aws.String(mfa)
		if p.MfaToken == "" {
			token, err := prompt(fmt.Sprintf("Enter token for %s: ", mfa))
			if err != nil {
				return sts.Credentials{}, err
			}
			params.TokenCode = aws.String(token)
		} else {
			params.TokenCode = p.MfaToken
		}<commit_msg>MFA token code must be cast to aws.String<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"time"

	"github.com/99designs/aws-vault/keyring"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/sts"
)

const (
	MaxSessionDuration    = time.Hour * 36
	MinSessionDuration    = time.Minute * 15
	MinAssumeRoleDuration = time.Minute * 15
	MaxAssumeRoleDuration = time.Hour

	DefaultSessionDuration    = time.Hour * 4
	DefaultAssumeRoleDuration = time.Minute * 15
)

type stsClient interface {
	AssumeRole(input *sts.AssumeRoleInput) (*sts.AssumeRoleOutput, error)
	GetSessionToken(input *sts.GetSessionTokenInput) (*sts.GetSessionTokenOutput, error)
}

type VaultOptions struct {
	SessionDuration    time.Duration
	AssumeRoleDuration time.Duration
	ExpiryWindow       time.Duration
	WriteEnv           bool
	MfaToken           string
}

func (o VaultOptions) Validate() error {
	if o.SessionDuration < time.Minute*15 {
		return errors.New("Minimum session duration is 15 minutes")
	} else if o.SessionDuration > time.Hour*36 {
		return errors.New("Maximum session duration is 36 hours")
	}
	return nil
}

func (o VaultOptions) ApplyDefaults() VaultOptions {
	if o.AssumeRoleDuration == 0 {
		o.AssumeRoleDuration = DefaultAssumeRoleDuration
	}
	if o.SessionDuration == 0 {
		o.SessionDuration = DefaultSessionDuration
	}
	return o
}

type VaultProvider struct {
	credentials.Expiry
	VaultOptions
	profile  string
	expires  time.Time
	keyring  keyring.Keyring
	profiles profiles
	session  *sts.Credentials
	client   stsClient
}

func NewVaultProvider(k keyring.Keyring, profile string, opts VaultOptions) (*VaultProvider, error) {
	opts = opts.ApplyDefaults()
	if err := opts.Validate(); err != nil {
		return nil, err
	}
	profiles, err := parseProfiles()
	if err != nil {
		return nil, err
	}
	return &VaultProvider{
		VaultOptions: opts,
		keyring:      k,
		profile:      profile,
		profiles:     profiles,
	}, nil
}

func (p *VaultProvider) Retrieve() (credentials.Value, error) {
	creds, err := p.getMasterCreds()
	if err != nil {
		return credentials.Value{}, err
	}

	session, err := p.getCachedSession()
	if err != nil {
		if err == keyring.ErrKeyNotFound {
			log.Println("Session not found in keyring")
		} else {
			log.Println(err)
		}

		session, err = p.getSessionToken(&creds)
		if err != nil {
			return credentials.Value{}, err
		}

		bytes, err := json.Marshal(session)
		if err != nil {
			return credentials.Value{}, err
		}

		log.Printf("Writing session for %s to keyring", p.profile)
		p.keyring.Set(keyring.Item{
			Key:       sessionKey(p.profile),
			Label:     "aws-vault session for " + p.profile,
			Data:      bytes,
			TrustSelf: true,
		})
	}

	log.Printf("Using session, expires in %s", session.Expiration.Sub(time.Now()).String())

	window := p.ExpiryWindow
	if window == 0 {
		window = p.SessionDuration - (p.SessionDuration / 3)
	}

	p.SetExpiration(*session.Expiration, window)
	p.expires = *session.Expiration

	if role, ok := p.profiles[p.profile]["role_arn"]; ok {
		session, err = p.assumeRole(session, role)
		if err != nil {
			return credentials.Value{}, err
		}

		log.Printf("Role token expires in %s", session.Expiration.Sub(time.Now()))
	}

	value := credentials.Value{
		AccessKeyID:     *session.AccessKeyId,
		SecretAccessKey: *session.SecretAccessKey,
		SessionToken:    *session.SessionToken,
	}

	return value, nil
}

func sessionKey(profile string) string {
	return profile + " session"
}

func (p *VaultProvider) getCachedSession() (session sts.Credentials, err error) {
	item, err := p.keyring.Get(sessionKey(p.profile))
	if err != nil {
		return session, err
	}

	if err = json.Unmarshal(item.Data, &session); err != nil {
		return session, err
	}

	if session.Expiration.Before(time.Now()) {
		return session, errors.New("Session is expired")
	}

	return
}

func (p *VaultProvider) getMasterCreds() (credentials.Value, error) {
	source := p.profiles.sourceProfile(p.profile)

	provider := credentials.NewChainCredentials([]credentials.Provider{
		&credentials.EnvProvider{},
		&credentials.SharedCredentialsProvider{Filename: "", Profile: p.profile},
		&KeyringProvider{Keyring: p.keyring, Profile: source},
	})

	return provider.Get()
}

func (p *VaultProvider) getSessionToken(creds *credentials.Value) (sts.Credentials, error) {
	params := &sts.GetSessionTokenInput{
		DurationSeconds: aws.Int64(int64(p.SessionDuration.Seconds())),
	}

	if mfa, ok := p.profiles[p.profile]["mfa_serial"]; ok {
		params.SerialNumber = aws.String(mfa)
		if p.MfaToken == "" {
			token, err := prompt(fmt.Sprintf("Enter token for %s: ", mfa))
			if err != nil {
				return sts.Credentials{}, err
			}
			params.TokenCode = aws.String(token)
		} else {
			params.TokenCode = aws.String(p.MfaToken)
		}
{'fraction_non_alphanumeric': 0.12687188019966722, 'fraction_numerical': 0.005615640599001664, 'mean_word_length': 3.9419321685508737, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5340875', 'n_tokens_mistral': 3496, 'n_tokens_neox': 3258, 'n_words': 753}