<commit_before>package vaulted

import (
	"errors"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/miquella/ask"
)

var STSDurationDefault = time.Hour

var (
	ErrInvalidCommand = errors.New("Invalid command")
)

type Vault struct {
	Vars     map[string]string `json:"vars"`
	AWSKey   *AWSKey           `json:"aws_key,omitempty"`
	SSHKeys  map[string]string `json:"ssh_keys,omitempty"`
	Duration time.Duration     `json:"duration,omitempty"`
}

type AWSKey struct {
	ID                      string `json:"id"`
	Secret                  string `json:"secret"`
	MFA                     string `json:"mfa,omitempty"`
	Role                    string `json:"role,omitempty"`
	ForgoTempCredGeneration bool   `json:"forgoTempCredGeneration"`
}

func (v *Vault) CreateEnvironment(extraVars map[string]string) (*Environment, error) {
	var duration time.Duration
	if v.Duration == 0 {
		duration = STSDurationDefault
	} else {
		duration = v.Duration
	}

	e := &Environment{
		Vars:       make(map[string]string),
		Expiration: time.Now().Add(duration).Unix(),
	}

	// copy the vault vars to the environment
	for key, value := range v.Vars {
		e.Vars[key] = value
	}
	for key, value := range extraVars {
		e.Vars[key] = value
	}

	// copy the vault ssh keys to the environment
	if len(v.SSHKeys) > 0 {
		e.SSHKeys = make(map[string]string)
		for key, value := range v.SSHKeys {
			e.SSHKeys[key] = value
		}
	}

	// get aws creds
	if v.AWSKey != nil && v.AWSKey.ID != "" && v.AWSKey.Secret != "" {
		var err error
		var creds *sts.Credentials
		if v.AWSKey.ForgoTempCredGeneration {
			creds = &sts.Credentials{
				AccessKeyId:     aws.String(v.AWSKey.ID),
				SecretAccessKey: aws.String(v.AWSKey.Secret),
				SessionToken:    aws.String(""), // permanent key, no session token
			}
		} else {
			if v.AWSKey.Role != "" {
				creds, err = v.AWSKey.assumeRole(duration)
			} else {
				creds, err = v.AWSKey.generateSTS(duration)
			}
		}

		if err != nil {
			return nil, err
		}

		e.Vars["AWS_ACCESS_KEY_ID"] = *creds.AccessKeyId
		e.Vars["AWS_SECRET_ACCESS_KEY"] = *creds.SecretAccessKey
		e.Vars["AWS_SESSION_TOKEN"] = *creds.SessionToken
		e.Vars["AWS_SECURITY_TOKEN"] = *creds.SessionToken
	}

	return e, nil
}

func (k *AWSKey) stsClient() *sts.STS {
	sess := session.New(&aws.Config{
		Credentials: credentials.NewStaticCredentials(
			k.ID,
			k.Secret,
			"", // Temporary session token
		),
	})
	return sts.New(sess)
}

func (k *AWSKey) getAssumeRoleInput(duration time.Duration) (*sts.AssumeRoleInput, error) {
	roleSessionName := "VaultedSession"
	input := &sts.AssumeRoleInput{
		DurationSeconds: aws.Int64(int64(duration.Seconds())),
		RoleArn:         &k.Role,
		RoleSessionName: &roleSessionName,
	}

	return input, nil
}

func (k *AWSKey) assumeRole(duration time.Duration) (*sts.Credentials, error) {
	// first generate a session token
	creds, err := k.generateSTS(duration)
	if err != nil {
		return nil, err
	}

	// now use the generated session token to assume the role
	sess := session.New(&aws.Config{
		Credentials: credentials.NewStaticCredentials(
			*creds.AccessKeyId,
			*creds.SecretAccessKey,
			*creds.SessionToken,
		),
	})
	stsClient := sts.New(sess)

	assumeRoleInput, err := k.getAssumeRoleInput(duration)
	if err != nil {
		return nil, err
	}

	resp, err := stsClient.AssumeRole(assumeRoleInput)
	if err != nil {
		return nil, err
	}

	return resp.Credentials, nil
}

func (k *AWSKey) buildSessionTokenInput(duration time.Duration) (*sts.GetSessionTokenInput, error) {
	input := &sts.GetSessionTokenInput{
		DurationSeconds: aws.Int64(int64(duration.Seconds())),
	}

	if k.MFA != "" {
		tokenCode, err := getTokenCode()
		if err != nil {
			return nil, err
		}<commit_msg>Use the caller identity to generate session names<commit_after>package vaulted

import (
	"errors"
	"fmt"
	"path"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/miquella/ask"
)

const (
	DefaultSessionName = "VaultedSession"
)

var STSDurationDefault = time.Hour

var (
	ErrInvalidCommand = errors.New("Invalid command")
)

type Vault struct {
	Vars     map[string]string `json:"vars"`
	AWSKey   *AWSKey           `json:"aws_key,omitempty"`
	SSHKeys  map[string]string `json:"ssh_keys,omitempty"`
	Duration time.Duration     `json:"duration,omitempty"`
}

type AWSKey struct {
	ID                      string `json:"id"`
	Secret                  string `json:"secret"`
	MFA                     string `json:"mfa,omitempty"`
	Role                    string `json:"role,omitempty"`
	ForgoTempCredGeneration bool   `json:"forgoTempCredGeneration"`
}

func (v *Vault) CreateEnvironment(extraVars map[string]string) (*Environment, error) {
	var duration time.Duration
	if v.Duration == 0 {
		duration = STSDurationDefault
	} else {
		duration = v.Duration
	}

	e := &Environment{
		Vars:       make(map[string]string),
		Expiration: time.Now().Add(duration).Unix(),
	}

	// copy the vault vars to the environment
	for key, value := range v.Vars {
		e.Vars[key] = value
	}
	for key, value := range extraVars {
		e.Vars[key] = value
	}

	// copy the vault ssh keys to the environment
	if len(v.SSHKeys) > 0 {
		e.SSHKeys = make(map[string]string)
		for key, value := range v.SSHKeys {
			e.SSHKeys[key] = value
		}
	}

	// get aws creds
	if v.AWSKey != nil && v.AWSKey.ID != "" && v.AWSKey.Secret != "" {
		var err error
		var creds *sts.Credentials
		if v.AWSKey.ForgoTempCredGeneration {
			creds = &sts.Credentials{
				AccessKeyId:     aws.String(v.AWSKey.ID),
				SecretAccessKey: aws.String(v.AWSKey.Secret),
				SessionToken:    aws.String(""), // permanent key, no session token
			}
		} else {
			if v.AWSKey.Role != "" {
				creds, err = v.AWSKey.assumeRole(duration)
			} else {
				creds, err = v.AWSKey.generateSTS(duration)
			}
		}

		if err != nil {
			return nil, err
		}

		e.Vars["AWS_ACCESS_KEY_ID"] = *creds.AccessKeyId
		e.Vars["AWS_SECRET_ACCESS_KEY"] = *creds.SecretAccessKey
		e.Vars["AWS_SESSION_TOKEN"] = *creds.SessionToken
		e.Vars["AWS_SECURITY_TOKEN"] = *creds.SessionToken
	}

	return e, nil
}

func (k *AWSKey) stsClient() *sts.STS {
	sess := session.New(&aws.Config{
		Credentials: credentials.NewStaticCredentials(
			k.ID,
			k.Secret,
			"", // Temporary session token
		),
	})
	return sts.New(sess)
}

func (k *AWSKey) assumeRole(duration time.Duration) (*sts.Credentials, error) {
	// first generate a session token
	creds, err := k.generateSTS(duration)
	if err != nil {
		return nil, err
	}

	// now use the generated session token to assume the role
	sess := session.New(&aws.Config{
		Credentials: credentials.NewStaticCredentials(
			*creds.AccessKeyId,
			*creds.SecretAccessKey,
			*creds.SessionToken,
		),
	})

	client := sts.New(sess)
	return k.assumeRoleWithClient(client, duration)
}

func (k *AWSKey) assumeRoleWithClient(client *sts.STS, duration time.Duration) (*sts.Credentials, error) {
	roleSessionName := DefaultSessionName

	callerIdentity, err := client.GetCallerIdentity(&sts.GetCallerIdentityInput{})
	if err == nil {
		parts := strings.SplitN(*callerIdentity.Arn, ":", 6)
		if len(parts) == 6 {
			roleSessionName = fmt.Sprintf("%s@%s", path.Base(parts[5]), parts[4])
		}
	}

	assumeRoleInput := &sts.AssumeRoleInput{
		DurationSeconds: aws.Int64(int64(duration.Seconds())),
		RoleArn:         &k.Role,
		RoleSessionName: &roleSessionName,
	}

	assumeRoleOutput, err := client.AssumeRole(assumeRoleInput)
	if err != nil {
		return nil, err
	}

	return assumeRoleOutput.Credentials, nil
}

func (k *AWSKey) buildSessionTokenInput(duration time.Duration) (*sts.GetSessionTokenInput, error) {
	input := &sts.GetSessionTokenInput{
		DurationSeconds: aws.Int64(int64(duration.Seconds())),
	}

	if k.MFA != "" {
		tokenCode, err := getTokenCode()
		if err != nil {
			return nil, err
		}
{'fraction_non_alphanumeric': 0.13901345291479822, 'fraction_numerical': 0.0029895366218236174, 'mean_word_length': 3.531038374717833, 'pattern_counts': {'":': 1, '<': 3, '<?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': '7076214', 'n_tokens_mistral': 3075, 'n_tokens_neox': 2894, 'n_words': 603}