Repository: pyca/cryptography (Python)

Issue:
[WIP] feat: add kdf_rounds param to encryption algorithm
### Summary

Currently the number of `kdf_rounds` used when encrypting OpenSSH private keys is fixed at the `ssh-keygen` default of `16`. `ssh-keygen` however parameterizes this value with the `-a` flag. This PR permits the same parameterization through the `BestAvailableEncryption` class.

### Open Questions
- [] Should this parameter be passed through `BestAvailableEncryption` where it is generally not applicable or should another data structure e.g. `KDFOptions` be created and passed through `BestAvailableEncryption` to better indicate this is optional and specific behavior?
- [] It doesn't really make sense to mention the default of `16` at the level of `BestAvailableEncryption` because it is only applied in the `ssh` module, but I don't have a better way to express that without discussing interface boundaries.
- [] Not sure which version this would be a candidate for so `version_added` is missing from the docs.

Requested originally in https://github.com/ansible-collections/community.crypto/issues/449


Hints:
Thank you for working on this! As you've discovered, our API doesn't have a clear place to override any options about how encryption is performed. Unfortunately, simply adding `kdf_rounds` to `BestAvailableEncryption` is probably not correct. We'll need to design an API that considers that KDF rounds may not be meaningful for some encryption options, and that it will have different meaning for different algorithms.

If you're interested in doing some of that design work, that's fantastic and we're very much up for working through these issues with you. https://github.com/pyca/cryptography/issues/4272 and https://github.com/pyca/cryptography/pull/5656 and https://github.com/pyca/cryptography/pull/6569 contain some of the background on the challenges here -- though we don't need to solve _all_ the problems discussed in them!
Thanks for the initial review. I'll put some thought into this, but my gut feeling is that OpenSSH private key encryption is not really aligned with the concept of `BestAvailableEncryption` to begin with i.e. it made reuse of this pattern for programming convenience and not for the sake of a coherent API. That being said to avoid breaking changes it is probably best to not touch that class for the time being so existing behavior remains intact and instead create a new encryption subclass closer to the OpenSSH domain. It may be a little clunky, but ultimately I think it's better than treating OpenSSH like it's OpenSSL.
Pushed a rough draft of a new `KeySerializationEncryption` interface which attempts to sprinkle in some polymorphism making new behavior less risky to implement. It does add some unnecessary detail to the abstract interface which could be pulled into a second abstraction layer like `PasswordBasedEncryption`, but not sure if that additional complexity is valuable at this point. The `options` property is a little loosely typed for my tastes, but then again adding a new data structure to contain the optional parameters isn't particularly 'pythonic'. Defaulted `kwargs` should avoid breaks and users shouldn't really instantiate `KeySerializationEncryption` directly, but it's public so a second set of eyes on that would  be nice. Choose not to update docs or address cov report until this draft is reviewed.

Patch:
diff --git a/src/cryptography/hazmat/backends/openssl/backend.py b/src/cryptography/hazmat/backends/openssl/backend.py
--- a/src/cryptography/hazmat/backends/openssl/backend.py
+++ b/src/cryptography/hazmat/backends/openssl/backend.py
@@ -1528,6 +1528,15 @@ def _private_key_bytes(
                     "Passwords longer than 1023 bytes are not supported by "
                     "this backend"
                 )
+        elif (
+            isinstance(
+                encryption_algorithm, serialization._KeySerializationEncryption
+            )
+            and encryption_algorithm._format
+            is format
+            is serialization.PrivateFormat.OpenSSH
+        ):
+            password = encryption_algorithm.password
         else:
             raise ValueError("Unsupported encryption type")
 
@@ -1592,7 +1601,9 @@ def _private_key_bytes(
         # OpenSSH + PEM
         if format is serialization.PrivateFormat.OpenSSH:
             if encoding is serialization.Encoding.PEM:
-                return ssh.serialize_ssh_private_key(key, password)
+                return ssh._serialize_ssh_private_key(
+                    key, password, encryption_algorithm
+                )
 
             raise ValueError(
                 "OpenSSH private key format can only be used"
diff --git a/src/cryptography/hazmat/primitives/_serialization.py b/src/cryptography/hazmat/primitives/_serialization.py
--- a/src/cryptography/hazmat/primitives/_serialization.py
+++ b/src/cryptography/hazmat/primitives/_serialization.py
@@ -3,6 +3,7 @@
 # for complete details.
 
 import abc
+import typing
 
 from cryptography import utils
 
@@ -25,6 +26,13 @@ class PrivateFormat(utils.Enum):
     Raw = "Raw"
     OpenSSH = "OpenSSH"
 
+    def encryption_builder(self) -> "KeySerializationEncryptionBuilder":
+        if self is not PrivateFormat.OpenSSH:
+            raise ValueError(
+                "encryption_builder only supported with PrivateFormat.OpenSSH"
+            )
+        return KeySerializationEncryptionBuilder(self)
+
 
 class PublicFormat(utils.Enum):
     SubjectPublicKeyInfo = "X.509 subjectPublicKeyInfo with PKCS#1"
@@ -53,3 +61,44 @@ def __init__(self, password: bytes):
 
 class NoEncryption(KeySerializationEncryption):
     pass
+
+
+class KeySerializationEncryptionBuilder(object):
+    def __init__(
+        self,
+        format: PrivateFormat,
+        *,
+        _kdf_rounds: typing.Optional[int] = None,
+    ) -> None:
+        self._format = format
+
+        self._kdf_rounds = _kdf_rounds
+
+    def kdf_rounds(self, rounds: int) -> "KeySerializationEncryptionBuilder":
+        if self._kdf_rounds is not None:
+            raise ValueError("kdf_rounds already set")
+        return KeySerializationEncryptionBuilder(
+            self._format, _kdf_rounds=rounds
+        )
+
+    def build(self, password: bytes) -> KeySerializationEncryption:
+        if not isinstance(password, bytes) or len(password) == 0:
+            raise ValueError("Password must be 1 or more bytes.")
+
+        return _KeySerializationEncryption(
+            self._format, password, kdf_rounds=self._kdf_rounds
+        )
+
+
+class _KeySerializationEncryption(KeySerializationEncryption):
+    def __init__(
+        self,
+        format: PrivateFormat,
+        password: bytes,
+        *,
+        kdf_rounds: typing.Optional[int],
+    ):
+        self._format = format
+        self.password = password
+
+        self._kdf_rounds = kdf_rounds
diff --git a/src/cryptography/hazmat/primitives/serialization/__init__.py b/src/cryptography/hazmat/primitives/serialization/__init__.py
--- a/src/cryptography/hazmat/primitives/serialization/__init__.py
+++ b/src/cryptography/hazmat/primitives/serialization/__init__.py
@@ -11,6 +11,7 @@
     ParameterFormat,
     PrivateFormat,
     PublicFormat,
+    _KeySerializationEncryption,
 )
 from cryptography.hazmat.primitives.serialization.base import (
     load_der_parameters,
@@ -42,4 +43,5 @@
     "KeySerializationEncryption",
     "BestAvailableEncryption",
     "NoEncryption",
+    "_KeySerializationEncryption",
 ]
diff --git a/src/cryptography/hazmat/primitives/serialization/ssh.py b/src/cryptography/hazmat/primitives/serialization/ssh.py
--- a/src/cryptography/hazmat/primitives/serialization/ssh.py
+++ b/src/cryptography/hazmat/primitives/serialization/ssh.py
@@ -15,9 +15,11 @@
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
 from cryptography.hazmat.primitives.serialization import (
     Encoding,
+    KeySerializationEncryption,
     NoEncryption,
     PrivateFormat,
     PublicFormat,
+    _KeySerializationEncryption,
 )
 
 try:
@@ -601,13 +603,13 @@ def load_ssh_private_key(
     return private_key
 
 
-def serialize_ssh_private_key(
+def _serialize_ssh_private_key(
     private_key: _SSH_PRIVATE_KEY_TYPES,
-    password: typing.Optional[bytes] = None,
+    password: bytes,
+    encryption_algorithm: KeySerializationEncryption,
 ) -> bytes:
     """Serialize private key with OpenSSH custom encoding."""
-    if password is not None:
-        utils._check_bytes("password", password)
+    utils._check_bytes("password", password)
 
     if isinstance(private_key, ec.EllipticCurvePrivateKey):
         key_type = _ecdsa_key_type(private_key.public_key())
@@ -628,6 +630,11 @@ def serialize_ssh_private_key(
         blklen = _SSH_CIPHERS[ciphername][3]
         kdfname = _BCRYPT
         rounds = _DEFAULT_ROUNDS
+        if (
+            isinstance(encryption_algorithm, _KeySerializationEncryption)
+            and encryption_algorithm._kdf_rounds is not None
+        ):
+            rounds = encryption_algorithm._kdf_rounds
         salt = os.urandom(16)
         f_kdfoptions.put_sshstr(salt)
         f_kdfoptions.put_u32(rounds)


Test patch:
diff --git a/tests/hazmat/primitives/test_serialization.py b/tests/hazmat/primitives/test_serialization.py
--- a/tests/hazmat/primitives/test_serialization.py
+++ b/tests/hazmat/primitives/test_serialization.py
@@ -2335,9 +2335,10 @@ def test_serialize_ssh_private_key_errors(self, backend):
 
         # bad object type
         with pytest.raises(ValueError):
-            ssh.serialize_ssh_private_key(
+            ssh._serialize_ssh_private_key(
                 object(),  # type:ignore[arg-type]
-                None,
+                b"",
+                NoEncryption(),
             )
 
         private_key = ec.generate_private_key(ec.SECP256R1(), backend)
@@ -2362,11 +2363,26 @@ def test_serialize_ssh_private_key_errors(self, backend):
             b"x" * 100,
         ),
     )
-    def test_serialize_ssh_private_key_with_password(self, password, backend):
+    @pytest.mark.parametrize(
+        "kdf_rounds",
+        [
+            1,
+            10,
+            30,
+        ],
+    )
+    def test_serialize_ssh_private_key_with_password(
+        self, password, kdf_rounds, backend
+    ):
         original_key = ec.generate_private_key(ec.SECP256R1(), backend)
-        encoded_key_data = ssh.serialize_ssh_private_key(
-            private_key=original_key,
-            password=password,
+        encoded_key_data = original_key.private_bytes(
+            Encoding.PEM,
+            PrivateFormat.OpenSSH,
+            (
+                PrivateFormat.OpenSSH.encryption_builder()
+                .kdf_rounds(kdf_rounds)
+                .build(password)
+            ),
         )
 
         decoded_key = load_ssh_private_key(
@@ -2416,3 +2432,22 @@ def test_dsa_private_key_sizes(self, key_path, supported, backend):
                 key.private_bytes(
                     Encoding.PEM, PrivateFormat.OpenSSH, NoEncryption()
                 )
+
+
+class TestEncryptionBuilder:
+    def test_unsupported_format(self):
+        f = PrivateFormat.PKCS8
+        with pytest.raises(ValueError):
+            f.encryption_builder()
+
+    def test_duplicate_kdf_rounds(self):
+        b = PrivateFormat.OpenSSH.encryption_builder().kdf_rounds(12)
+        with pytest.raises(ValueError):
+            b.kdf_rounds(12)
+
+    def test_invalid_password(self):
+        b = PrivateFormat.OpenSSH.encryption_builder()
+        with pytest.raises(ValueError):
+            b.build(12)  # type: ignore[arg-type]
+        with pytest.raises(ValueError):
+            b.build(b"")