<commit_before>package vault

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/sso"
	"github.com/aws/aws-sdk-go/service/ssooidc"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/skratchdot/open-golang/open"
)

type OIDCTokenCacher interface {
	Get(string) (*ssooidc.CreateTokenOutput, error)
	Set(string, *ssooidc.CreateTokenOutput) error
}

// SSORoleCredentialsProvider creates temporary credentials for an SSO Role.
type SSORoleCredentialsProvider struct {
	OIDCClient     *ssooidc.SSOOIDC
	OIDCTokenCache OIDCTokenCacher
	StartURL       string
	SSOClient      *sso.SSO
	AccountID      string
	RoleName       string
	ExpiryWindow   time.Duration
	credentials.Expiry
}

// Retrieve generates a new set of temporary credentials using SSO GetRoleCredentials.
func (p *SSORoleCredentialsProvider) Retrieve() (credentials.Value, error) {
	creds, err := p.getRoleCredentials()
	if err != nil {
		return credentials.Value{}, err
	}

	p.SetExpiration(aws.MillisecondsTimeValue(creds.Expiration), p.ExpiryWindow)
	return credentials.Value{
		AccessKeyID:     aws.StringValue(creds.AccessKeyId),
		SecretAccessKey: aws.StringValue(creds.SecretAccessKey),
		SessionToken:    aws.StringValue(creds.SessionToken),
	}, nil
}

func (p *SSORoleCredentialsProvider) getRoleCredentials() (*sso.RoleCredentials, error) {
	token, err := p.getOIDCToken()
	if err != nil {
		return nil, err
	}

	resp, err := p.SSOClient.GetRoleCredentials(&sso.GetRoleCredentialsInput{
		AccessToken: token.AccessToken,
		AccountId:   aws.String(p.AccountID),
		RoleName:    aws.String(p.RoleName),
	})
	if err != nil {
		return nil, err
	}
	log.Printf("Got credentials %s for SSO role %s (account: %s), expires in %s", FormatKeyForDisplay(*resp.RoleCredentials.AccessKeyId), p.RoleName, p.AccountID, time.Until(aws.MillisecondsTimeValue(resp.RoleCredentials.Expiration)).String())

	return resp.RoleCredentials, nil
}

// getRoleCredentialsAsStsCredemtials returns getRoleCredentials as sts.Credentials because sessions.Store expects it
func (p *SSORoleCredentialsProvider) getRoleCredentialsAsStsCredemtials() (*sts.Credentials, error) {
	creds, err := p.getRoleCredentials()
	if err != nil {
		return nil, err
	}

	return &sts.Credentials{
		AccessKeyId:     creds.AccessKeyId,
		SecretAccessKey: creds.SecretAccessKey,
		SessionToken:    creds.SessionToken,
		Expiration:      aws.Time(aws.MillisecondsTimeValue(creds.Expiration)),
	}, nil
}

func (p *SSORoleCredentialsProvider) getOIDCToken() (token *ssooidc.CreateTokenOutput, err error) {
	if p.OIDCTokenCache != nil {
		token, err = p.OIDCTokenCache.Get(p.StartURL)
		if err != nil {
			return nil, err
		}
	}
	if token == nil {
		token, err = p.newOIDCToken()
		if err != nil {
			return nil, err
		}

		if p.OIDCTokenCache != nil {
			err = p.OIDCTokenCache.Set(p.StartURL, token)
			return nil, err
		}
	}
	return token, err
}
<commit_msg>Handle OIDC caching for new keyring entries<commit_after>package vault

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/99designs/keyring"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/sso"
	"github.com/aws/aws-sdk-go/service/ssooidc"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/skratchdot/open-golang/open"
)

type OIDCTokenCacher interface {
	Get(string) (*ssooidc.CreateTokenOutput, error)
	Set(string, *ssooidc.CreateTokenOutput) error
}

// SSORoleCredentialsProvider creates temporary credentials for an SSO Role.
type SSORoleCredentialsProvider struct {
	OIDCClient     *ssooidc.SSOOIDC
	OIDCTokenCache OIDCTokenCacher
	StartURL       string
	SSOClient      *sso.SSO
	AccountID      string
	RoleName       string
	ExpiryWindow   time.Duration
	credentials.Expiry
}

// Retrieve generates a new set of temporary credentials using SSO GetRoleCredentials.
func (p *SSORoleCredentialsProvider) Retrieve() (credentials.Value, error) {
	creds, err := p.getRoleCredentials()
	if err != nil {
		return credentials.Value{}, err
	}

	p.SetExpiration(aws.MillisecondsTimeValue(creds.Expiration), p.ExpiryWindow)
	return credentials.Value{
		AccessKeyID:     aws.StringValue(creds.AccessKeyId),
		SecretAccessKey: aws.StringValue(creds.SecretAccessKey),
		SessionToken:    aws.StringValue(creds.SessionToken),
	}, nil
}

func (p *SSORoleCredentialsProvider) getRoleCredentials() (*sso.RoleCredentials, error) {
	token, err := p.getOIDCToken()
	if err != nil {
		return nil, err
	}

	resp, err := p.SSOClient.GetRoleCredentials(&sso.GetRoleCredentialsInput{
		AccessToken: token.AccessToken,
		AccountId:   aws.String(p.AccountID),
		RoleName:    aws.String(p.RoleName),
	})
	if err != nil {
		return nil, err
	}
	log.Printf("Got credentials %s for SSO role %s (account: %s), expires in %s", FormatKeyForDisplay(*resp.RoleCredentials.AccessKeyId), p.RoleName, p.AccountID, time.Until(aws.MillisecondsTimeValue(resp.RoleCredentials.Expiration)).String())

	return resp.RoleCredentials, nil
}

// getRoleCredentialsAsStsCredemtials returns getRoleCredentials as sts.Credentials because sessions.Store expects it
func (p *SSORoleCredentialsProvider) getRoleCredentialsAsStsCredemtials() (*sts.Credentials, error) {
	creds, err := p.getRoleCredentials()
	if err != nil {
		return nil, err
	}

	return &sts.Credentials{
		AccessKeyId:     creds.AccessKeyId,
		SecretAccessKey: creds.SecretAccessKey,
		SessionToken:    creds.SessionToken,
		Expiration:      aws.Time(aws.MillisecondsTimeValue(creds.Expiration)),
	}, nil
}

func (p *SSORoleCredentialsProvider) getOIDCToken() (token *ssooidc.CreateTokenOutput, err error) {
	if p.OIDCTokenCache != nil {
		token, err = p.OIDCTokenCache.Get(p.StartURL)
		if err != nil && err != keyring.ErrKeyNotFound {
			return nil, err
		}
	}
	if token == nil {
		token, err = p.newOIDCToken()
		if err != nil {
			return nil, err
		}

		err = p.OIDCTokenCache.Set(p.StartURL, token)
		if err != nil {
			return nil, err
		}
	}
	return token, err
}

{'fraction_non_alphanumeric': 0.11769520379767556, 'fraction_numerical': 0.0003273858241938124, 'mean_word_length': 4.470008952551477, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '506560', 'n_tokens_mistral': 2226, 'n_tokens_neox': 2173, 'n_words': 437}