Support RSA_SIGN_PSS* GCP keys

src/pykmstool/kms_funcs.py

@@ -30,7 +30,11 @@ def kms_sign_csr(*, client: kms.KeyManagementServiceClient, key_version_name: st
         .subject_name(name)
     )
 
-    cert = builder.sign(signer_priv_key, signer_priv_key.hash_algorithm())
+    cert = builder.sign(
+        private_key=signer_priv_key,
+        algorithm=signer_priv_key.hash_algorithm(),
+        rsa_padding=signer_priv_key.rsa_padding(signer_priv_key.hash_algorithm())
+    )
     return cert.public_bytes(serialization.Encoding.PEM).decode('ascii')
 
 

src/pykmstool/kms_priv_key/__init__.py

@@ -7,6 +7,7 @@
 
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurve
+from cryptography.hazmat.primitives.asymmetric.padding import PKCS1v15, PSS, MGF1
 from cryptography.hazmat.primitives.hashes import SHA256, SHA384, SHA512
 from google.cloud.kms_v1 import KeyManagementServiceClient, CryptoKeyVersion
 
@@ -17,14 +18,18 @@
 
 def build_kms_priv_key(
         cls: typing.Type[KMSRSAPrivateKey | KMSECPrivateKey | KMSEd25519PrivateKey],
-        hash_algorithm: typing.Type[SHA256 | SHA384 | SHA512] = None,
-        curve: typing.Type[EllipticCurve] = None,
+        hash_algorithm: typing.Optional[typing.Type[SHA256 | SHA384 | SHA512]] = None,
+        rsa_padding: typing.Optional[typing.Callable[[SHA256 | SHA384 | SHA512], PKCS1v15 | PSS]] = None,
+        curve: typing.Optional[typing.Type[EllipticCurve]] = None,
 ) -> typing.Callable[..., KMSRSAPrivateKey | KMSECPrivateKey | KMSEd25519PrivateKey]:
     bind_kwargs = {}
 
     if hash_algorithm:
         bind_kwargs.update({"hash_algorithm": hash_algorithm})
 
+    if rsa_padding:
+        bind_kwargs.update({"rsa_padding": rsa_padding})
+
     if curve:
         bind_kwargs.update({"curve": curve})
 
@@ -33,15 +38,27 @@ def build_kms_priv_key(
 
 def create_pyca_private_key(client: KeyManagementServiceClient, key_version_name: str)\
         -> KMSRSAPrivateKey | KMSECPrivateKey | KMSEd25519PrivateKey:
+    f_pkcs1v15 = lambda _: PKCS1v15()
+    f_pss = lambda hash_alg: PSS(MGF1(hash_alg), PSS.DIGEST_LENGTH)
+
     kms_alg_to_class = {
         CryptoKeyVersion.CryptoKeyVersionAlgorithm.RSA_SIGN_PKCS1_2048_SHA256.name:
-            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA256),
+            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA256, rsa_padding=f_pkcs1v15),
         CryptoKeyVersion.CryptoKeyVersionAlgorithm.RSA_SIGN_PKCS1_3072_SHA256.name:
-            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA256),
+            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA256, rsa_padding=f_pkcs1v15),
         CryptoKeyVersion.CryptoKeyVersionAlgorithm.RSA_SIGN_PKCS1_4096_SHA256.name:
-            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA256),
+            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA256, rsa_padding=f_pkcs1v15),
         CryptoKeyVersion.CryptoKeyVersionAlgorithm.RSA_SIGN_PKCS1_4096_SHA512.name:
-            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA512),
+            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA512, rsa_padding=f_pkcs1v15),
+
+        CryptoKeyVersion.CryptoKeyVersionAlgorithm.RSA_SIGN_PSS_2048_SHA256.name:
+            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA256, rsa_padding=f_pss),
+        CryptoKeyVersion.CryptoKeyVersionAlgorithm.RSA_SIGN_PSS_3072_SHA256.name:
+            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA256, rsa_padding=f_pss),
+        CryptoKeyVersion.CryptoKeyVersionAlgorithm.RSA_SIGN_PSS_4096_SHA256.name:
+            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA256, rsa_padding=f_pss),
+        CryptoKeyVersion.CryptoKeyVersionAlgorithm.RSA_SIGN_PSS_4096_SHA512.name:
+            build_kms_priv_key(KMSRSAPrivateKey, hash_algorithm=SHA512, rsa_padding=f_pss),
 
         CryptoKeyVersion.CryptoKeyVersionAlgorithm.EC_SIGN_P256_SHA256.name:
             build_kms_priv_key(KMSECPrivateKey, hash_algorithm=SHA256, curve=ec.SECP256R1),

src/pykmstool/kms_priv_key/base_key.py

@@ -11,13 +11,12 @@
 
 from cryptography import utils
 from cryptography.hazmat.primitives import hashes
-from cryptography.hazmat.primitives.asymmetric import (
-    utils as asym_utils,
-)
+from cryptography.hazmat.primitives.asymmetric.padding import PSS, PKCS1v15
 from cryptography.hazmat.primitives.hashes import SHA256, SHA384, SHA512
 from cryptography.hazmat.primitives.serialization import load_pem_public_key
 from google.cloud.kms_v1 import KeyManagementServiceClient, CryptoKeyVersion
 
+
 def crc32c(data: bytes) -> int:
     import crcmod  # type: ignore
 
@@ -30,20 +29,26 @@ def __init__(
             self,
             client: KeyManagementServiceClient,
             ckv: CryptoKeyVersion,
-            hash_algorithm: typing.Type[SHA256 | SHA384 | SHA512] | typing.Callable[[], None]
+            hash_algorithm: typing.Type[SHA256 | SHA384 | SHA512] | typing.Callable[[], None],
+            rsa_padding: typing.Callable[[SHA256 | SHA384 | SHA512 | None], PKCS1v15 | PSS | None],
     ):
         self._client = client
         self._ckv = ckv
         self._hash_algorithm = hash_algorithm
+        self._rsa_padding = rsa_padding
 
     @property
     def crypto_key_version(self):
         return self._ckv
 
     @property
-    def hash_algorithm(self) -> typing.Type[SHA256 | SHA384 | SHA512]:
+    def hash_algorithm(self) -> typing.Type[SHA256 | SHA384 | SHA512] | typing.Callable[[], None]:
         return self._hash_algorithm
 
+    @property
+    def rsa_padding(self) -> typing.Callable[[SHA256 | SHA384 | SHA512 | None], PKCS1v15 | PSS | None]:
+        return self._rsa_padding
+
     def _common_public_key(self):
         public_key = self._client.get_public_key(name=self._ckv.name)
 
@@ -52,33 +57,32 @@ def _common_public_key(self):
 
         return load_pem_public_key(public_key.pem.encode("ascii"))
 
-    def _common_sign(self, data: utils.Buffer, algorithm: hashes.HashAlgorithm | typing.Callable[[], None]) -> bytes:
-        if isinstance(algorithm, asym_utils.Prehashed):
-            raise RuntimeError("Prehashed data is not supported.")
+    def _common_sign(self, data: utils.Buffer, algorithm: hashes.HashAlgorithm | None = None) -> bytes:
+        instance_hash_obj = self.hash_algorithm()
 
-        if not self.hash_algorithm():
-            if algorithm():
+        if not instance_hash_obj:
+            if algorithm:
                 raise RuntimeError("Unexpected algorithm parameter provided.")
 
             sign_response = self._client.asymmetric_sign(
                 request={
                     "name": self._ckv.name,
                     "data": base64.b64encode(data).decode('ascii'),
-                    "data_crc32c": crc32c(data),
+                    "data_crc32c": crc32c(bytes(data)),
                 }
             )
         else:
-            if algorithm.name != self.hash_algorithm.name:
-               raise RuntimeError("Requested incompatible hash algorithm.")
+            if algorithm and algorithm.name != instance_hash_obj.name:
+                raise RuntimeError("Requested incompatible hash algorithm.")
 
-            h = hashes.Hash(self.hash_algorithm())
+            h = hashes.Hash(instance_hash_obj)
             h.update(data)
             digest = h.finalize()
 
             sign_response = self._client.asymmetric_sign(
                 request={
                     "name": self._ckv.name,
-                    "digest": {self.hash_algorithm.name: base64.b64encode(digest).decode('ascii')},
+                    "digest": {instance_hash_obj.name: base64.b64encode(digest).decode('ascii')},
                     "digest_crc32c": crc32c(digest),
                 }
             )
@@ -87,4 +91,3 @@ def _common_sign(self, data: utils.Buffer, algorithm: hashes.HashAlgorithm | typ
             raise RuntimeError("Mismatched CRC32C in the signature returned from KMS.")
 
         return sign_response.signature
-

src/pykmstool/kms_priv_key/key_ec.py

@@ -13,6 +13,7 @@
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurveSignatureAlgorithm, EllipticCurve, ECDH, \
     EllipticCurvePublicKey, EllipticCurvePrivateKey, EllipticCurvePrivateNumbers
+from cryptography.hazmat.primitives.asymmetric.utils import Prehashed
 from cryptography.hazmat.primitives.hashes import SHA256, SHA384, SHA512
 from google.cloud.kms_v1 import KeyManagementServiceClient, CryptoKeyVersion
 
@@ -27,26 +28,25 @@ def __init__(
             hash_algorithm: typing.Type[SHA256 | SHA384 | SHA512],
             curve: typing.Type[ec.EllipticCurve]
     ):
-        super().__init__(client, ckv, hash_algorithm)
+        super().__init__(client, ckv, hash_algorithm, rsa_padding=lambda _: None)
         self._curve = curve
 
     def __copy__(self) -> EllipticCurvePrivateKey:
         return KMSECPrivateKey(
             client=self._client,
             ckv=self._ckv,
-            hash_algorithm=self._hash_algorithm,
+            hash_algorithm=self._hash_algorithm,  # noqa
             curve=self._curve
         )
 
     def public_key(self) -> EllipticCurvePublicKey:
         return self._common_public_key()
 
     def sign(self, data: utils.Buffer, signature_algorithm: EllipticCurveSignatureAlgorithm) -> bytes:
-        return self._common_sign(data, signature_algorithm.algorithm)
+        if isinstance(signature_algorithm.algorithm, Prehashed):
+            raise RuntimeError("Prehashed data is not supported.")
 
-    @property
-    def hash_algorithm(self) -> typing.Type[SHA256 | SHA384 | SHA512]:
-        return self._hash_algorithm
+        return self._common_sign(data, signature_algorithm.algorithm)  # noqa
 
     @property
     def curve(self) -> EllipticCurve:

src/pykmstool/kms_priv_key/key_ed25519.py

@@ -21,7 +21,7 @@ def __init__(
             client: KeyManagementServiceClient,
             ckv: CryptoKeyVersion
     ):
-        super().__init__(client, ckv, hash_algorithm=lambda: None)
+        super().__init__(client, ckv, hash_algorithm=lambda: None, rsa_padding=lambda _: None)
 
     def __copy__(self) -> Ed25519PrivateKey:
         return KMSEd25519PrivateKey(
@@ -33,7 +33,7 @@ def public_key(self) -> Ed25519PublicKey:
         return self._common_public_key()
 
     def sign(self, data: Buffer) -> bytes:
-        return self._common_sign(data, algorithm=lambda: None)
+        return self._common_sign(data)
 
     def private_bytes(self, encoding: serialization.Encoding, format: serialization.PrivateFormat,
                       encryption_algorithm: serialization.KeySerializationEncryption) -> bytes:

src/pykmstool/kms_priv_key/key_rsa.py

@@ -9,11 +9,8 @@
 import typing
 
 from cryptography.hazmat.primitives import serialization, hashes
-from cryptography.hazmat.primitives.asymmetric import (
-    rsa,
-    utils as asym_utils,
-)
-from cryptography.hazmat.primitives.asymmetric.padding import AsymmetricPadding
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.primitives.asymmetric.padding import AsymmetricPadding, PSS
 from cryptography.hazmat.primitives.asymmetric.padding import PKCS1v15
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey
 from cryptography.hazmat.primitives.hashes import SHA256, SHA384, SHA512
@@ -23,14 +20,26 @@
 
 
 class KMSRSAPrivateKey(rsa.RSAPrivateKey, BaseKMSPrivateKey):
-    def __init__(self, client: KeyManagementServiceClient, ckv: CryptoKeyVersion, hash_algorithm: typing.Type[SHA256 | SHA384 | SHA512]):
-        super().__init__(client, ckv, hash_algorithm)
+    def __init__(
+            self,
+            client: KeyManagementServiceClient,
+            ckv: CryptoKeyVersion,
+            hash_algorithm: typing.Type[SHA256 | SHA384 | SHA512],
+            rsa_padding: typing.Callable[[SHA256 | SHA384 | SHA512], PKCS1v15 | PSS],
+    ):
+        super().__init__(
+            client=client,
+            ckv=ckv,
+            hash_algorithm=hash_algorithm,
+            rsa_padding=rsa_padding  # noqa
+        )
 
     def __copy__(self) -> RSAPrivateKey:
         return KMSRSAPrivateKey(
             client=self._client,
             ckv=self._ckv,
-            hash_algorithm=self._hash_algorithm
+            hash_algorithm=self._hash_algorithm,  # noqa
+            rsa_padding=self._rsa_padding  # noqa
         )
 
     def public_key(self) -> RSAPublicKey:
@@ -40,13 +49,14 @@ def sign(
         self,
         data: bytes,
         padding: AsymmetricPadding,
-        algorithm: typing.Union[asym_utils.Prehashed, hashes.HashAlgorithm],
+        algorithm: hashes.HashAlgorithm,
     ) -> bytes:
-        if not isinstance(padding, PKCS1v15):
+        if not isinstance(padding, PKCS1v15) and not isinstance(padding, PSS):
             raise RuntimeError("Unsupported padding type requested.")
 
         return self._common_sign(data, algorithm)
 
+    @property
     def key_size(self) -> int:
         return self.public_key().key_size