Merge pull request #15 from cormacdalton/fix-encryption

Fix encryption and decryption
Mastercard · Jan 31, 2024 · 44c0e1d · 44c0e1d
2 parents 2350cc7 + 536bdfe
commit 44c0e1d
Show file tree

Hide file tree

Showing 6 changed files with 188 additions and 10 deletions.
diff --git a/encryption/encryption.go b/encryption/encryption.go
@@ -3,10 +3,14 @@ package encryption
 import (
 	"github.com/Jeffail/gabs/v2"
 	"github.com/mastercard/client-encryption-go/jwe"
+	"github.com/mastercard/client-encryption-go/utils"
 )
 
 func EncryptPayload(payload string, config jwe.JWEConfig) string {
-	jsonPayload, _ := gabs.ParseJSON([]byte(payload))
+	jsonPayload, err := gabs.ParseJSON([]byte(payload))
+	if err != nil {
+		panic(err)
+	}
 	for jsonPathIn, jsonPathOut := range config.GetEncryptionPaths() {
 		jsonPayload = encryptPayloadPath(jsonPayload, jsonPathIn, jsonPathOut, config)
 	}
@@ -28,24 +32,31 @@ func encryptPayloadPath(jsonPayload *gabs.Container, jsonPathIn string, jsonPath
 		Kid: config.GetEncryptionKeyFingerprint(),
 		Cty: "application/json",
 	}
-
-	payload, err := jwe.Encrypt(config, jsonPayload.String(), joseHeader)
+	jsonPathIn = utils.RemoveRoot(jsonPathIn)
+	jsonPathOut = utils.RemoveRoot(jsonPathOut)
+	payloadToEncrypt := utils.GetPayloadToEncrypt(jsonPayload, jsonPathIn)
+	payload, err := jwe.Encrypt(config, payloadToEncrypt, joseHeader)
 	if err != nil {
 		panic(err)
 	}
-
 	if jsonPathIn == "$" {
 		jsonPayload = gabs.New()
 	} else {
 		jsonPayload.DeleteP(jsonPathIn)
 	}
-	jsonPayload.Set(payload, config.GetEncryptedValueFieldName())
+	if jsonPathOut == "$" {
+		jsonPayload.SetP(payload, config.GetEncryptedValueFieldName())
+	} else {
+		jsonPayload.SetP(payload, jsonPathOut+"."+config.GetEncryptedValueFieldName())
+	}
 	return jsonPayload
 }
 
 func decryptPayloadPath(jsonPayload *gabs.Container, jsonPathIn string, jsonPathOut string, config jwe.JWEConfig) *gabs.Container {
-	inJsonObject := jsonPayload.Path(jsonPathIn).Data().(string)
-	jweObject, err := jwe.ParseJWEObject(inJsonObject)
+	jsonPathIn = utils.RemoveRoot(jsonPathIn)
+	jsonPathOut = utils.RemoveRoot(jsonPathOut)
+	encryptedPayload := utils.GetPayloadToDecrypt(jsonPayload, jsonPathIn)
+	jweObject, err := jwe.ParseJWEObject(encryptedPayload)
 	if err != nil {
 		panic(err)
 	}
@@ -57,9 +68,8 @@ func decryptPayloadPath(jsonPayload *gabs.Container, jsonPathIn string, jsonPath
 	if jsonPathOut == "$" {
 		jsonPayload = jsonDecryptedPayload
 	} else {
-		jsonPayload.DeleteP(jsonPathIn)
-		jsonPayload.Set(jsonDecryptedPayload, jsonPathOut)
+		jsonPayload.DeleteP(utils.GetParent(jsonPathIn))
+		jsonPayload.SetP(jsonDecryptedPayload, jsonPathOut)
 	}
-
 	return jsonPayload
 }
diff --git a/encryption/encryption_test.go b/encryption/encryption_test.go
@@ -37,6 +37,31 @@ func TestEncryptPayload_ShouldEncryptRootArrays(t *testing.T) {
 	assert.Equal(t, "[{},{}]", decryptedPayload)
 }
 
+func TestEncryptPayload_ShouldDecryptWholePayload(t *testing.T) {
+	const payload = `{"value1":{"value2":10,"value3":20},"value4":30}`
+
+	decryptionKeyPath := "../testdata/keys/pkcs8/test_key_pkcs8-2048.der"
+	certificatePath := "../testdata/certificates/test_certificate-2048.der"
+
+	decryptionKey, err := utils.LoadUnencryptedDecryptionKey(decryptionKeyPath)
+	assert.Nil(t, err)
+	certificate, err := utils.LoadEncryptionCertificate(certificatePath)
+	assert.Nil(t, err)
+
+	cb := jwe.NewJWEConfigBuilder()
+	jweConfig := cb.WithDecryptionKey(decryptionKey).
+		WithEncryptionPath("$", "$").
+		WithDecryptionPath("$", "$").
+		WithEncryptedValueFieldName("encryptedData").
+		WithCertificate(certificate).
+		Build()
+
+	encryptedPayload := encryption.EncryptPayload(payload, *jweConfig)
+	assert.Regexp(t, `{"encryptedData":"[^"]*?"}`, encryptedPayload)
+	decryptedPayload := encryption.DecryptPayload(encryptedPayload, *jweConfig)
+	assert.Equal(t, payload, decryptedPayload)
+}
+
 func TestDecryptPayload_ShouldEncryptRootArrays(t *testing.T) {
 	const encryptedPayload = "{" +
 		"    \"encryptedData\": \"eyJraWQiOiI3NjFiMDAzYzFlYWRlM2E1NDkwZTUwMDBkMzc4ODdiYWE1ZTZlYzBlMjI2YzA3NzA2ZTU5OTQ1MWZjMDMyYTc5IiwiY3R5IjoiYXBwbGljYXRpb24vanNvbiIsImVuYyI6IkEyNTZHQ00iLCJhbGciOiJSU0EtT0FFUC0yNTYifQ.IcTIce59pgtjODJn4PhR7oK3F-gxcd7dishTrT7T9y5VC0U5ZS_JdMoRe59_UTkJMY8Nykb2rv3Oh_jSDYRmGB_CWMIciXYMLHQptLTF5xI1ZauDPnooDMWoOCBD_d3I0wTJNcM7I658rK0ZWSByVK9YqhEo8UaIf4e6egRHQdZ2_IGKgICwmglv_uXQrYewOWFTKR1uMpya1N50MDnWax2NtnW3SljP3mARUBLBnRmOyubQCg-Mgn8fsOWWXm-KL9RrQq9AF_HJceoJl1rRgzPW7g6SLK6EjiGW_ArTmrLaOHg9bYOY_LrbyokK_M1pMo9qup70DHvjHkMZqIL3aQ.vtma3jBIo2STkquxTUX9PQ.9ZoQG0sFvQ.ms4bW3OFd03neRlex-zZ8w\"" +
@@ -60,3 +85,53 @@ func TestDecryptPayload_ShouldEncryptRootArrays(t *testing.T) {
 	decryptedPayload := encryption.DecryptPayload(encryptedPayload, *jweConfig)
 	assert.Equal(t, "[{},{}]", decryptedPayload)
 }
+
+func TestEncryptPayload_ShouldEncryptExampleFromREADME(t *testing.T) {
+	const payload = `{"path":{"to":{"foo":{"sensitiveField1":"sensitiveValue1","sensitiveField2":"sensitiveValue2"}}}}`
+
+	decryptionKeyPath := "../testdata/keys/pkcs8/test_key_pkcs8-2048.der"
+	certificatePath := "../testdata/certificates/test_certificate-2048.der"
+
+	decryptionKey, err := utils.LoadUnencryptedDecryptionKey(decryptionKeyPath)
+	assert.Nil(t, err)
+	certificate, err := utils.LoadEncryptionCertificate(certificatePath)
+	assert.Nil(t, err)
+
+	cb := jwe.NewJWEConfigBuilder()
+	jweConfig := cb.WithDecryptionKey(decryptionKey).
+		WithEncryptionPath("$.path.to.foo", "$.path.to.encryptedFoo").
+		WithDecryptionPath("$.path.to.encryptedFoo.encryptedData", "$.path.to.foo").
+		WithEncryptedValueFieldName("encryptedData").
+		WithCertificate(certificate).
+		Build()
+
+	encryptedPayload := encryption.EncryptPayload(payload, *jweConfig)
+	assert.Regexp(t, `{"path":{"to":{"encryptedFoo":{"encryptedData":"[^"]*?"}}}}`, encryptedPayload)
+	decryptedPayload := encryption.DecryptPayload(encryptedPayload, *jweConfig)
+	assert.Equal(t, payload, decryptedPayload)
+}
+
+func TestEncryptPayload_ShouldEncryptWithoutRootKey(t *testing.T) {
+	const payload = `{"path":{"to":{"foo":{"sensitiveField1":"sensitiveValue1","sensitiveField2":"sensitiveValue2"}}}}`
+
+	decryptionKeyPath := "../testdata/keys/pkcs8/test_key_pkcs8-2048.der"
+	certificatePath := "../testdata/certificates/test_certificate-2048.der"
+
+	decryptionKey, err := utils.LoadUnencryptedDecryptionKey(decryptionKeyPath)
+	assert.Nil(t, err)
+	certificate, err := utils.LoadEncryptionCertificate(certificatePath)
+	assert.Nil(t, err)
+
+	cb := jwe.NewJWEConfigBuilder()
+	jweConfig := cb.WithDecryptionKey(decryptionKey).
+		WithEncryptionPath("path.to.foo", "path.to.encryptedFoo").
+		WithDecryptionPath("path.to.encryptedFoo.encryptedData", "path.to.foo").
+		WithEncryptedValueFieldName("encryptedData").
+		WithCertificate(certificate).
+		Build()
+
+	encryptedPayload := encryption.EncryptPayload(payload, *jweConfig)
+	assert.Regexp(t, `{"path":{"to":{"encryptedFoo":{"encryptedData":"[^"]*?"}}}}`, encryptedPayload)
+	decryptedPayload := encryption.DecryptPayload(encryptedPayload, *jweConfig)
+	assert.Equal(t, payload, decryptedPayload)
+}
diff --git a/utils/encryption_utils.go b/utils/encryption_utils.go
@@ -4,6 +4,7 @@ import (
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/pem"
+	"github.com/Jeffail/gabs/v2"
 	"golang.org/x/crypto/pkcs12"
 	"io/ioutil"
 	"os"
@@ -73,6 +74,25 @@ func LoadUnencryptedDecryptionKey(keyFilePath string) (*rsa.PrivateKey, error) {
 	return key.(*rsa.PrivateKey), nil
 }
 
+func GetPayloadToEncrypt(jsonPayload *gabs.Container, jsonPathIn string) string {
+	if jsonPathIn == "$" {
+		return jsonPayload.String()
+	} else {
+		return jsonPayload.Path(jsonPathIn).String()
+	}
+}
+
+func GetPayloadToDecrypt(jsonPayload *gabs.Container, jsonPathIn string) string {
+	if jsonPathIn == "$" {
+		return jsonPayload.
+			Children()[0].
+			Data().(string)
+	} else {
+		return jsonPayload.Path(jsonPathIn).
+			Data().(string)
+	}
+}
+
 // Read File
 func readFile(path string) ([]byte, error) {
 	file, err := os.Open(path)

diff --git a/utils/encryption_utils_test.go b/utils/encryption_utils_test.go
@@ -1,6 +1,7 @@
 package utils_test
 
 import (
+	"github.com/Jeffail/gabs/v2"
 	"testing"
 
 	"github.com/mastercard/client-encryption-go/utils"
@@ -30,3 +31,35 @@ func TestLoadUnencryptedDecryptionKey(t *testing.T) {
 	_, err := utils.LoadUnencryptedDecryptionKey(keyPath)
 	assert.Nil(t, err)
 }
+
+func TestGetPayloadToEncryptWithRootPayload_ShouldReturnEntirePayload(t *testing.T) {
+	jsonPath := "$"
+	payload := `{"path":{"to":{"foo":{"sensitiveField1":"sensitiveValue1","sensitiveField2":"sensitiveValue2"}}}}`
+	jsonPayload, _ := gabs.ParseJSON([]byte(payload))
+	payloadToEncrypt := utils.GetPayloadToEncrypt(jsonPayload, jsonPath)
+	assert.Equal(t, payload, payloadToEncrypt)
+}
+
+func TestGetPayloadToEncryptWithPayloadPath_ShouldReturnPayloadAtPath(t *testing.T) {
+	jsonPath := "path.to.foo"
+	payload := `{"path":{"to":{"foo":{"sensitiveField1":"sensitiveValue1","sensitiveField2":"sensitiveValue2"}}}}`
+	jsonPayload, _ := gabs.ParseJSON([]byte(payload))
+	payloadToEncrypt := utils.GetPayloadToEncrypt(jsonPayload, jsonPath)
+	assert.Equal(t, `{"sensitiveField1":"sensitiveValue1","sensitiveField2":"sensitiveValue2"}`, payloadToEncrypt)
+}
+
+func TestGetPayloadToDecryptWithRootPayload_ShouldReturnEntirePayload(t *testing.T) {
+	jsonPath := "$"
+	payload := `{"encryptedPayload":"abcdefg"}`
+	jsonPayload, _ := gabs.ParseJSON([]byte(payload))
+	payloadToDecrypt := utils.GetPayloadToDecrypt(jsonPayload, jsonPath)
+	assert.Equal(t, "abcdefg", payloadToDecrypt)
+}
+
+func TestGetPayloadToDecryptWithRootPayload_ShouldReturnPayloadAtPath(t *testing.T) {
+	jsonPath := "path.to.encryptedFoo.encryptedData"
+	payload := `{"path":{"to":{"encryptedFoo":{"encryptedData":"abcdefg"}}}}`
+	jsonPayload, _ := gabs.ParseJSON([]byte(payload))
+	payloadToDecrypt := utils.GetPayloadToDecrypt(jsonPayload, jsonPath)
+	assert.Equal(t, "abcdefg", payloadToDecrypt)
+}
diff --git a/utils/json_utils.go b/utils/json_utils.go
@@ -0,0 +1,16 @@
+package utils
+
+import "strings"
+
+func RemoveRoot(json string) string {
+	if json[0] == '$' && json != "$" {
+		return json[2:]
+	}
+	return json
+}
+
+func GetParent(path string) string {
+	keys := strings.Split(path, ".")
+	parent := keys[:len(keys)-1]
+	return strings.Join(parent, ".")
+}
diff --git a/utils/json_utils_test.go b/utils/json_utils_test.go
@@ -0,0 +1,24 @@
+package utils
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestRemoveRoot_ShouldRemoveRootIfPresent(t *testing.T) {
+	path := "$.path.to.foo"
+	updatedPath := RemoveRoot(path)
+	assert.Equal(t, "path.to.foo", updatedPath)
+}
+
+func TestRemoveRoot_ShouldNotRemoveRootIfItIsARootPath(t *testing.T) {
+	path := "$"
+	updatedPath := RemoveRoot(path)
+	assert.Equal(t, "$", updatedPath)
+}
+
+func TestGetParent(t *testing.T) {
+	path := "$.path.to.foo"
+	parent := GetParent(path)
+	assert.Equal(t, "$.path.to", parent)
+}