1package s3crypto
2
3import (
4	"github.com/aws/aws-sdk-go/aws"
5	"github.com/aws/aws-sdk-go/aws/awserr"
6	"github.com/aws/aws-sdk-go/service/kms"
7	"github.com/aws/aws-sdk-go/service/kms/kmsiface"
8)
9
10const (
11	// KMSWrap is a constant used during decryption to build a KMS key handler.
12	KMSWrap = "kms"
13)
14
15// kmsKeyHandler will make calls to KMS to get the masterkey
16type kmsKeyHandler struct {
17	kms   kmsiface.KMSAPI
18	cmkID *string
19
20	CipherData
21}
22
23// NewKMSKeyGenerator builds a new KMS key provider using the customer key ID and material
24// description.
25//
26// Example:
27//	sess := session.New(&aws.Config{})
28//	cmkID := "arn to key"
29//	matdesc := s3crypto.MaterialDescription{}
30//	handler := s3crypto.NewKMSKeyGenerator(kms.New(sess), cmkID)
31func NewKMSKeyGenerator(kmsClient kmsiface.KMSAPI, cmkID string) CipherDataGenerator {
32	return NewKMSKeyGeneratorWithMatDesc(kmsClient, cmkID, MaterialDescription{})
33}
34
35// NewKMSKeyGeneratorWithMatDesc builds a new KMS key provider using the customer key ID and material
36// description.
37//
38// Example:
39//	sess := session.New(&aws.Config{})
40//	cmkID := "arn to key"
41//	matdesc := s3crypto.MaterialDescription{}
42//	handler := s3crypto.NewKMSKeyGeneratorWithMatDesc(kms.New(sess), cmkID, matdesc)
43func NewKMSKeyGeneratorWithMatDesc(kmsClient kmsiface.KMSAPI, cmkID string, matdesc MaterialDescription) CipherDataGenerator {
44	if matdesc == nil {
45		matdesc = MaterialDescription{}
46	}
47	matdesc["kms_cmk_id"] = &cmkID
48
49	// These values are read only making them thread safe
50	kp := &kmsKeyHandler{
51		kms:   kmsClient,
52		cmkID: &cmkID,
53	}
54	// These values are read only making them thread safe
55	kp.CipherData.WrapAlgorithm = KMSWrap
56	kp.CipherData.MaterialDescription = matdesc
57	return kp
58}
59
60// NewKMSWrapEntry builds returns a new KMS key provider and its decrypt handler.
61//
62// Example:
63//	sess := session.New(&aws.Config{})
64//	customKMSClient := kms.New(sess)
65//	decryptHandler := s3crypto.NewKMSWrapEntry(customKMSClient)
66//
67//	svc := s3crypto.NewDecryptionClient(sess, func(svc *s3crypto.DecryptionClient) {
68//		svc.WrapRegistry[s3crypto.KMSWrap] = decryptHandler
69//	}))
70func NewKMSWrapEntry(kmsClient kmsiface.KMSAPI) WrapEntry {
71	// These values are read only making them thread safe
72	kp := &kmsKeyHandler{
73		kms: kmsClient,
74	}
75
76	return kp.decryptHandler
77}
78
79// decryptHandler initializes a KMS keyprovider with a material description. This
80// is used with Decrypting kms content, due to the cmkID being in the material description.
81func (kp kmsKeyHandler) decryptHandler(env Envelope) (CipherDataDecrypter, error) {
82	m := MaterialDescription{}
83	err := m.decodeDescription([]byte(env.MatDesc))
84	if err != nil {
85		return nil, err
86	}
87
88	cmkID, ok := m["kms_cmk_id"]
89	if !ok {
90		return nil, awserr.New("MissingCMKIDError", "Material description is missing CMK ID", nil)
91	}
92
93	kp.CipherData.MaterialDescription = m
94	kp.cmkID = cmkID
95	kp.WrapAlgorithm = KMSWrap
96	return &kp, nil
97}
98
99// DecryptKey makes a call to KMS to decrypt the key.
100func (kp *kmsKeyHandler) DecryptKey(key []byte) ([]byte, error) {
101	return kp.DecryptKeyWithContext(aws.BackgroundContext(), key)
102}
103
104// DecryptKeyWithContext makes a call to KMS to decrypt the key with request context.
105func (kp *kmsKeyHandler) DecryptKeyWithContext(ctx aws.Context, key []byte) ([]byte, error) {
106	out, err := kp.kms.DecryptWithContext(ctx,
107		&kms.DecryptInput{
108			EncryptionContext: kp.CipherData.MaterialDescription,
109			CiphertextBlob:    key,
110			GrantTokens:       []*string{},
111		})
112	if err != nil {
113		return nil, err
114	}
115	return out.Plaintext, nil
116}
117
118// GenerateCipherData makes a call to KMS to generate a data key, Upon making
119// the call, it also sets the encrypted key.
120func (kp *kmsKeyHandler) GenerateCipherData(keySize, ivSize int) (CipherData, error) {
121	return kp.GenerateCipherDataWithContext(aws.BackgroundContext(), keySize, ivSize)
122}
123
124// GenerateCipherDataWithContext makes a call to KMS to generate a data key,
125// Upon making the call, it also sets the encrypted key.
126func (kp *kmsKeyHandler) GenerateCipherDataWithContext(ctx aws.Context, keySize, ivSize int) (CipherData, error) {
127	out, err := kp.kms.GenerateDataKeyWithContext(ctx,
128		&kms.GenerateDataKeyInput{
129			EncryptionContext: kp.CipherData.MaterialDescription,
130			KeyId:             kp.cmkID,
131			KeySpec:           aws.String("AES_256"),
132		})
133	if err != nil {
134		return CipherData{}, err
135	}
136
137	iv := generateBytes(ivSize)
138	cd := CipherData{
139		Key:                 out.Plaintext,
140		IV:                  iv,
141		WrapAlgorithm:       KMSWrap,
142		MaterialDescription: kp.CipherData.MaterialDescription,
143		EncryptedKey:        out.CiphertextBlob,
144	}
145	return cd, nil
146}
147