Adding Unit Tests for BYOK Feature

1. add BYOK interface for python/rust/go sdk and add test
2. format python files
3. disabled timeout setting in Python sdk to pass BYOK testing

Signed-off-by: lisiteng <[email protected]>

sdk/go/client.go

@@ -2,7 +2,10 @@ package ehsm
 
 import (
 	"bytes"
+	"crypto"
 	hmac "crypto/hmac"
+	"crypto/rand"
+	"crypto/rsa"
 	sha256 "crypto/sha256"
 	tls "crypto/tls"
 	"encoding/base64"
@@ -139,3 +142,36 @@ func paramsSortStr(signParams *orderedmap.OrderedMap) string {
 	}
 	return str
 }
+
+func generateRandomKey(keyspec string) ([]byte, error) {
+	switch keyspec {
+	case "EH_AES_GCM_128", "EH_SM4_CBC", "EH_SM4_CTR":
+		key := make([]byte, 16)
+		return key, nil
+	case "EH_AES_GCM_192":
+		key := make([]byte, 24)
+		return key, nil
+	case "EH_AES_GCM_256":
+		key := make([]byte, 32)
+		return key, nil
+	default:
+		return nil, fmt.Errorf("Invalid keyspec")
+	}
+}
+
+func rsaEncrypt(pubkey crypto.PublicKey, key []byte, padding_mode string) ([]byte, error) {
+	if padding_mode == "EH_RSA_PKCS1" {
+		key_material, err := rsa.EncryptPKCS1v15(rand.Reader, pubkey.(*rsa.PublicKey), []byte(key))
+		if err != nil {
+			return nil, fmt.Errorf("EncryptPKCS1v15 failed.")
+		}
+		return key_material, nil
+	} else if padding_mode == "EH_RSA_PKCS1_OAEP" {
+		key_material, err := rsa.EncryptOAEP(sha256.New(), rand.Reader, pubkey.(*rsa.PublicKey), []byte(key), nil)
+		if err != nil {
+			return nil, fmt.Errorf("EncryptOAEP failed.")
+		}
+		return key_material, nil
+	}
+	return nil, fmt.Errorf("The padding mode is not supported.")
+}

sdk/go/get_parameters_for_import.go

@@ -0,0 +1,117 @@
+package ehsm
+
+import (
+	"crypto"
+	"crypto/x509"
+	"encoding/pem"
+	"errors"
+	"fmt"
+
+	"github.com/iancoleman/orderedmap"
+)
+
+/*
+Description:
+Queries the parameters that are used to import key material for a customer master key (CMK).
+Input:
+keyid  -- The unique keyid of import CMK.
+keyspec  -- The type of the public key that is used to encrypt key material.
+Output:
+pubkey  -- The public key that is used to encrypt key material.The public key is Base64-encoded.
+import_token  -- The import_token is Base64-encoded string.
+*/
+func (c *Client) GetParametersForImport(keyid, keyspec string) (crypto.PublicKey, string, error) {
+	payload := orderedmap.New()
+
+	payload.Set("keyid", keyid)
+	payload.Set("keyspec", keyspec)
+	params := c.initParams(payload)
+
+	c.modifyLock.Lock()
+	defer c.modifyLock.Unlock()
+
+	// call ehsm kms
+	resp, err := c.doPost(params, "GetParametersForImport")
+	if err != nil {
+		return "", "", err
+	}
+	result, ok := resp["result"].(map[string]interface{})
+
+	if !ok {
+		return "", "", fmt.Errorf("result field is not a valid map")
+	}
+	pubkey, ok := result["pubkey"].(string)
+	if !ok {
+		return "", "", fmt.Errorf("pubkey field is not a valid string")
+	}
+
+	import_token, ok := result["importToken"].(string)
+	if !ok {
+		return "", "", fmt.Errorf("import_token field is not a valid string")
+	}
+	pemBytes := []byte(pubkey)
+	derBytes, _ := pem.Decode(pemBytes)
+	if derBytes == nil {
+		return "", "", errors.New("PEM decoding failed")
+	}
+	switch derBytes.Type {
+	case string(PublicKeyPEMType):
+		pub_key, err := x509.ParsePKIXPublicKey(derBytes.Bytes)
+		if err != nil {
+			return "", "", err
+		}
+		return pub_key, import_token, nil
+	case string(PKCS1PublicKeyPEMType):
+		pub_key, err := x509.ParsePKCS1PublicKey(derBytes.Bytes)
+		if err != nil {
+			return "", "", err
+		}
+		return pub_key, import_token, nil
+	default:
+		return "", "", fmt.Errorf("unknown Public key PEM file type: %v. Are you passing the correct public key?",
+			derBytes.Type)
+	}
+
+}
+
+/*
+Description:
+Call the ImportKeyMaterial operation to import the key material.
+Input:
+keyid  -- The unique keyid of import CMK.
+padding_mode  -- The padding mode that is used to encrypt key material.
+key_material  -- The encrypted symmetric key.
+import_token  -- The import_token is Base64-encoded string.
+Output:
+message  -- The description of result（success or failed）.
+*/
+func (c *Client) ImportKeyMaterial(keyid, padding_mode, key_material, import_token string) (bool, error) {
+
+	payload := orderedmap.New()
+
+	payload.Set("keyid", keyid)
+	payload.Set("padding_mode", padding_mode)
+	payload.Set("key_material", key_material)
+	payload.Set("importToken", import_token)
+
+	params := c.initParams(payload)
+
+	c.modifyLock.Lock()
+	defer c.modifyLock.Unlock()
+
+	// call ehsm kms
+	resp, err := c.doPost(params, "ImportKeyMaterial")
+	if err != nil {
+		return false, err
+	}
+	resultString, ok := resp["result"].(map[string]interface{})
+	if !ok {
+		return false, fmt.Errorf("import_token field is not a valid string")
+	}
+	result, ok := resultString["result"].(bool)
+	if !ok {
+		return false, fmt.Errorf("result field is not a valid bool")
+	}
+
+	return result, nil
+}

sdk/go/getpublickey.go

@@ -34,7 +34,7 @@ func (c *Client) GetPublicKey(keyid string) (crypto.PublicKey, error) {
 	if keyid != "" {
 		payload.Set("keyid", keyid)
 	} else {
-		return "", fmt.Errorf("Please input keyid.")
+		return nil, fmt.Errorf("Please input keyid.")
 	}
 
 	params := c.initParams(payload)
@@ -45,16 +45,16 @@ func (c *Client) GetPublicKey(keyid string) (crypto.PublicKey, error) {
 	// call ehsm kms
 	resp, err := c.doPost(params, "GetPublicKey")
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	result, ok := resp["result"].(map[string]interface{})
 	if !ok {
-		return "", fmt.Errorf("result field is not a valid map")
+		return nil, fmt.Errorf("result field is not a valid map")
 	}
 
 	pubkey, ok := result["pubkey"].(string)
 	if !ok {
-		return "", fmt.Errorf("pubkey field is not a valid string")
+		return nil, fmt.Errorf("pubkey field is not a valid string")
 	}
 	// parse keyblob to crypto.x509
 	// support EH_RSA_x and EH_EC_P256

sdk/go/kms_test.go

@@ -1,6 +1,7 @@
 package ehsm
 
 import (
+	"crypto/rand"
 	"crypto/sha256"
 	"encoding/base64"
 	"testing"
@@ -186,3 +187,82 @@ func TestInvalidKeyUsage(t *testing.T) {
 		t.Error("keyusage mismatch test failed.")
 	}
 }
+
+func TestImportKey(t *testing.T) {
+	client, err := NewClient()
+	if err != nil {
+		t.Error(err)
+	}
+	keyspec := []string{"EH_AES_GCM_128", "EH_AES_GCM_192", "EH_AES_GCM_256", "EH_SM4_CBC", "EH_SM4_CTR"}
+	warpping_keyspec := []string{"EH_RSA_2048", "EH_RSA_3072", "EH_RSA_4096"}
+	padding_mode := []string{"EH_RSA_PKCS1", "EH_RSA_PKCS1_OAEP"}
+	for _, keyspec := range keyspec {
+		for _, warpping_keyspec := range warpping_keyspec {
+			for _, padding_mode := range padding_mode {
+				keyid, err := client.CreateKey(keyspec, "EH_EXTERNAL_KEY", "EH_KEYUSAGE_ENCRYPT_DECRYPT")
+				if err != nil {
+					t.Error(err)
+				}
+				pubkey, importToken, err := client.GetParametersForImport(keyid, warpping_keyspec)
+				if err != nil {
+					t.Error(err)
+				}
+				key, err := generateRandomKey(keyspec)
+				if err != nil {
+					t.Error(err)
+				}
+				if _, err := rand.Read(key); err != nil {
+					t.Error(err)
+				}
+				key_material, err := rsaEncrypt(pubkey, key, padding_mode)
+				if err != nil {
+					t.Error(err)
+				}
+
+				_, err = client.ImportKeyMaterial(keyid, padding_mode, base64.StdEncoding.EncodeToString(key_material), importToken)
+				if err != nil {
+					t.Error(err)
+				}
+				msg := "test for encrypt"
+				aad := "aad"
+				ciphertext, err := client.Encrypt(keyid, base64.StdEncoding.EncodeToString([]byte(msg)), base64.StdEncoding.EncodeToString([]byte(aad)))
+				if err != nil {
+					t.Error(err)
+				}
+				plaintext_b64, err := client.Decrypt(keyid, ciphertext, base64.StdEncoding.EncodeToString([]byte(aad)))
+				if err != nil {
+					t.Error(err)
+				}
+				plaintext, _ := base64.StdEncoding.DecodeString(plaintext_b64)
+				if err != nil || string(plaintext) != msg {
+					t.Error(err)
+				}
+			}
+		}
+	}
+}
+
+func TestEncryptDecrypt(t *testing.T) {
+	client, err := NewClient()
+	if err != nil {
+		t.Error(err)
+	}
+	keyid, err := client.CreateKey("EH_AES_GCM_256", "EH_INTERNAL_KEY", "EH_KEYUSAGE_ENCRYPT_DECRYPT")
+	if err != nil {
+		t.Error(err)
+	}
+	msg := "test for encrypt"
+	aad := "aad"
+	ciphertext, err := client.Encrypt(keyid, base64.StdEncoding.EncodeToString([]byte(msg)), base64.StdEncoding.EncodeToString([]byte(aad)))
+	if err != nil {
+		panic(err)
+	}
+	plaintext_b64, err := client.Decrypt(keyid, ciphertext, base64.StdEncoding.EncodeToString([]byte(aad)))
+	if err != nil {
+		panic(err)
+	}
+	plaintext, _ := base64.StdEncoding.DecodeString(plaintext_b64)
+	if err != nil || string(plaintext) != msg {
+		t.Error(err)
+	}
+}

sdk/python/ehsm/__init__.py

@@ -5,4 +5,4 @@
 from . import serializers
 from . import exceptions
 from . import session
-from . import utils
+from . import utils

sdk/python/ehsm/api/__init__.py

@@ -23,7 +23,7 @@ def __init__(
         session: Optional[Session] = None,
         appid: str = "",
         apikey: str = "",
-        allow_insecure: bool = False
+        allow_insecure: bool = False,
     ) -> None:
         """
         Initialize APIClient for EHSM KMS service

sdk/python/ehsm/api/crypto.py

@@ -138,3 +138,29 @@ def get_public_key(self, keyid: str):
             "", params={"Action": "GetPublicKey"}, data={"keyid": keyid}
         )
         return GetPublicKeyResponse.from_response(resp)
+
+    def import_key_material(
+        self, keyid: str, key_material: str, padding_mode: PaddingMode, importToken: str
+    ):
+        resp = self._session.post(
+            "",
+            params={"Action": "ImportKeyMaterial"},
+            data={
+                "keyid": keyid,
+                "key_material": key_material,
+                "padding_mode": padding_mode,
+                "importToken": importToken,
+            },
+        )
+        return ImportKeyMaterialResponse.from_response(resp)
+
+    def get_parameters_for_import(self, keyid: str, keyspec: KeySpec):
+        resp = self._session.post(
+            "",
+            params={"Action": "GetParametersForImport"},
+            data={
+                "keyid": keyid,
+                "keyspec": keyspec,
+            },
+        )
+        return GetParamtersForImportResponse.from_response(resp)