signing and encryption keys in gcp, aws, azure kms (#307)

* signing and encryption keys in gcp, aws, azure kms

* minor change

Author: rootflo-hardik

wavefront/server/packages/flo_cloud/flo_cloud/aws/kms.py

@@ -5,6 +5,7 @@
 from .._types import FloKMS
 
 aws_kms_arn = os.getenv('AWS_KMS_ARN')
+aws_kms_enc_arn = os.getenv('AWS_KMS_ENC_ARN')
 aws_region = os.getenv('AWS_REGION')
 
 
@@ -14,16 +15,23 @@ def __init__(self):
             raise ValueError('Region and KMS ARN must be set')
 
         self.aws_kms_arn = aws_kms_arn
+        self.aws_kms_enc_arn = aws_kms_enc_arn
         self.aws_region = aws_region
         self.kms_client = boto3.client('kms', region_name=aws_region)
 
     def encrypt(self, plaintext: str) -> bytes:
-        return self.kms_client.encrypt(KeyId=self.aws_kms_arn, Plaintext=plaintext)
+        if not self.aws_kms_enc_arn:
+            raise ValueError('AWS_KMS_ENC_ARN must be set to use encryption')
+        return self.kms_client.encrypt(KeyId=self.aws_kms_enc_arn, Plaintext=plaintext)[
+            'CiphertextBlob'
+        ]
 
     def decrypt(self, ciphertext: str) -> bytes:
+        if not self.aws_kms_enc_arn:
+            raise ValueError('AWS_KMS_ENC_ARN must be set to use decryption')
         return self.kms_client.decrypt(
-            KeyId=self.aws_kms_arn, CiphertextBlob=ciphertext
-        )
+            KeyId=self.aws_kms_enc_arn, CiphertextBlob=ciphertext
+        )['Plaintext']
 
     def sign(self, message: bytes, **kwargs) -> bytes:
         signing_algorithm = kwargs.get('signing_algorithm', 'RSASSA_PSS_SHA_256')

wavefront/server/packages/flo_cloud/flo_cloud/azure/key_vault.py

@@ -37,6 +37,8 @@ def __init__(
         vault_url: Optional[str] = None,
         key_name: Optional[str] = None,
         key_version: Optional[str] = None,
+        enc_key_name: Optional[str] = None,
+        enc_key_version: Optional[str] = None,
         client_id: Optional[str] = None,
         client_secret: Optional[str] = None,
         tenant_id: Optional[str] = None,
@@ -46,6 +48,12 @@ def __init__(
         resolved_key_version = key_version or os.environ.get(
             'AZURE_KEY_VAULT_KEY_VERSION'
         )
+        resolved_enc_key_name = enc_key_name or os.environ.get(
+            'AZURE_KEY_VAULT_ENC_KEY_NAME'
+        )
+        resolved_enc_key_version = enc_key_version or os.environ.get(
+            'AZURE_KEY_VAULT_ENC_KEY_VERSION'
+        )
 
         if not resolved_vault_url:
             raise ValueError(
@@ -74,19 +82,43 @@ def __init__(
         self._key_name = resolved_key_name
         self._key_version = resolved_key_version
         self.key_client = KeyClient(vault_url=resolved_vault_url, credential=credential)
-        key = self.key_client.get_key(resolved_key_name, version=resolved_key_version)
-        self.crypto_client = CryptographyClient(key, credential=credential)
+
+        sign_key = self.key_client.get_key(
+            resolved_key_name, version=resolved_key_version
+        )
+        self.crypto_client = CryptographyClient(sign_key, credential=credential)
+
+        self.enc_crypto_client = (
+            CryptographyClient(
+                self.key_client.get_key(
+                    resolved_enc_key_name, version=resolved_enc_key_version
+                ),
+                credential=credential,
+            )
+            if resolved_enc_key_name
+            else None
+        )
 
     def encrypt(self, plaintext: str) -> bytes:
+        if not self.enc_crypto_client:
+            raise ValueError(
+                'AZURE_KEY_VAULT_ENC_KEY_NAME must be set to use encryption'
+            )
         if isinstance(plaintext, str):
             plaintext = plaintext.encode('utf-8')
-        result = self.crypto_client.encrypt(EncryptionAlgorithm.rsa_oaep_256, plaintext)
+        result = self.enc_crypto_client.encrypt(
+            EncryptionAlgorithm.rsa_oaep_256, plaintext
+        )
         return result.ciphertext
 
     def decrypt(self, ciphertext: str) -> bytes:
+        if not self.enc_crypto_client:
+            raise ValueError(
+                'AZURE_KEY_VAULT_ENC_KEY_NAME must be set to use decryption'
+            )
         if isinstance(ciphertext, str):
             ciphertext = ciphertext.encode('utf-8')
-        result = self.crypto_client.decrypt(
+        result = self.enc_crypto_client.decrypt(
             EncryptionAlgorithm.rsa_oaep_256, ciphertext
         )
         return result.plaintext

wavefront/server/packages/flo_cloud/flo_cloud/gcp/kms.py

@@ -43,20 +43,20 @@ def __init__(self):
             crypto_key_version=gcp_crypto_key_version,
         )
 
-        # Separate symmetric key path for encryption/decryption.
-        # Symmetric encrypt/decrypt operates on a CryptoKey (not a version);
-        # GCP picks the primary version on encrypt and reads the embedded
-        # version metadata on decrypt.
-        enc_key = gcp_enc_crypto_key or gcp_crypto_key
-
-        self.enc_key_name = self.kms_client.crypto_key_path(
-            project=gcp_project_id,
-            location=gcp_location,
-            key_ring=gcp_key_ring,
-            crypto_key=enc_key,
+        self.enc_key_name = (
+            self.kms_client.crypto_key_path(
+                project=gcp_project_id,
+                location=gcp_location,
+                key_ring=gcp_key_ring,
+                crypto_key=gcp_enc_crypto_key,
+            )
+            if gcp_enc_crypto_key
+            else None
         )
 
     def encrypt(self, plaintext: bytes | str) -> bytes:
+        if not self.enc_key_name:
+            raise ValueError('GCP_KMS_ENC_CRYPTO_KEY must be set to use encryption')
         if isinstance(plaintext, str):
             plaintext = plaintext.encode('utf-8')
         request = kms_v1.EncryptRequest(
@@ -67,6 +67,8 @@ def encrypt(self, plaintext: bytes | str) -> bytes:
         return response.ciphertext
 
     def decrypt(self, ciphertext: bytes | str) -> bytes:
+        if not self.enc_key_name:
+            raise ValueError('GCP_KMS_ENC_CRYPTO_KEY must be set to use decryption')
         if isinstance(ciphertext, str):
             ciphertext = ciphertext.encode('utf-8')
         request = kms_v1.DecryptRequest(