feat(xtest): Adds pure mlkem test scenarios

Author: dmihalcik-virtru

xtest/abac.py

@@ -173,6 +173,8 @@ class KasGrantValue(BaseModelIgnoreExtra):
     "hpqt:xwing",
     "hpqt:secp256r1-mlkem768",
     "hpqt:secp384r1-mlkem1024",
+    "mlkem:768",
+    "mlkem:1024",
 ]
 
 KAS_PUBLIC_KEY_ALG_ENUM_RSA_2048 = 1
@@ -186,6 +188,10 @@ class KasGrantValue(BaseModelIgnoreExtra):
 KAS_PUBLIC_KEY_ALG_ENUM_HPQT_SECP256R1_MLKEM768 = 11
 KAS_PUBLIC_KEY_ALG_ENUM_HPQT_SECP384R1_MLKEM1024 = 12
 
+# Pure ML-KEM enums match protobuf ALGORITHM_MLKEM_768/1024 in platform PR #3537.
+KAS_PUBLIC_KEY_ALG_ENUM_MLKEM_768 = 20
+KAS_PUBLIC_KEY_ALG_ENUM_MLKEM_1024 = 21
+
 _KAS_ALG_TO_STR_MAP = {
     KAS_PUBLIC_KEY_ALG_ENUM_RSA_2048: "rsa:2048",
     KAS_PUBLIC_KEY_ALG_ENUM_RSA_4096: "rsa:4096",
@@ -195,6 +201,8 @@ class KasGrantValue(BaseModelIgnoreExtra):
     KAS_PUBLIC_KEY_ALG_ENUM_HPQT_XWING: "hpqt:xwing",
     KAS_PUBLIC_KEY_ALG_ENUM_HPQT_SECP256R1_MLKEM768: "hpqt:secp256r1-mlkem768",
     KAS_PUBLIC_KEY_ALG_ENUM_HPQT_SECP384R1_MLKEM1024: "hpqt:secp384r1-mlkem1024",
+    KAS_PUBLIC_KEY_ALG_ENUM_MLKEM_768: "mlkem:768",
+    KAS_PUBLIC_KEY_ALG_ENUM_MLKEM_1024: "mlkem:1024",
 }
 _STR_TO_KAS_ALG_MAP = {v: k for k, v in _KAS_ALG_TO_STR_MAP.items()}
 

xtest/conftest.py

@@ -173,12 +173,13 @@ def pytest_configure(config: pytest.Config) -> None:
     # If the user passed --sdks-encrypt / --sdks-decrypt explicitly, their
     # tokens win and we skip the resolution step entirely.
     need_resolve = (
-        (not config.getoption("--sdks-encrypt") and scenario.sdks.encrypt)
-        or (not config.getoption("--sdks-decrypt") and scenario.sdks.decrypt)
-    )
+        not config.getoption("--sdks-encrypt") and scenario.sdks.encrypt
+    ) or (not config.getoption("--sdks-decrypt") and scenario.sdks.decrypt)
     if need_resolve:
         try:
-            tokens = scenario_to_pytest_sdks(scenario, installed_json_for(scenario_path))
+            tokens = scenario_to_pytest_sdks(
+                scenario, installed_json_for(scenario_path)
+            )
         except FileNotFoundError as e:
             raise pytest.UsageError(str(e)) from e
         if not config.getoption("--sdks-encrypt") and tokens["encrypt"]:

xtest/fixtures/keys.py

@@ -261,6 +261,30 @@ def key_secpmlkem_5(
     )
 
 
+@pytest.fixture(scope="module")
+def key_mlkem_768(
+    otdfctl: OpentdfCommandLineTool,
+    kas_entry_km1: abac.KasEntry,
+    root_key: str,
+) -> abac.KasKey:
+    """Get or create pure ML-KEM-768 managed key 'm1' on km1."""
+    return _get_or_create_key(
+        otdfctl, kas_entry_km1, "m1", "mlkem:768", root_key, "mechanism-mlkem"
+    )
+
+
+@pytest.fixture(scope="module")
+def key_mlkem_1024(
+    otdfctl: OpentdfCommandLineTool,
+    kas_entry_km1: abac.KasEntry,
+    root_key: str,
+) -> abac.KasKey:
+    """Get or create pure ML-KEM-1024 managed key 'm2' on km1."""
+    return _get_or_create_key(
+        otdfctl, kas_entry_km1, "m2", "mlkem:1024", root_key, "mechanism-mlkem"
+    )
+
+
 # ---------------------------------------------------------------------------
 # Attribute + key assignment fixtures (value-level)
 # ---------------------------------------------------------------------------
@@ -385,6 +409,42 @@ def attribute_with_secpmlkem_5_key(
     )
 
 
+@pytest.fixture(scope="module")
+def attribute_with_mlkem_768_key(
+    otdfctl: OpentdfCommandLineTool,
+    key_mlkem_768: abac.KasKey,
+    otdf_client_scs: abac.SubjectConditionSet,
+    temporary_namespace: abac.Namespace,
+) -> tuple[abac.Attribute, list[str]]:
+    """ALL_OF attribute with one pure ML-KEM-768 key assigned at value level."""
+    return _create_keyed_attribute(
+        otdfctl,
+        temporary_namespace,
+        "mlkem768-test",
+        [("m1", key_mlkem_768)],
+        otdf_client_scs,
+        "mechanism-mlkem",
+    )
+
+
+@pytest.fixture(scope="module")
+def attribute_with_mlkem_1024_key(
+    otdfctl: OpentdfCommandLineTool,
+    key_mlkem_1024: abac.KasKey,
+    otdf_client_scs: abac.SubjectConditionSet,
+    temporary_namespace: abac.Namespace,
+) -> tuple[abac.Attribute, list[str]]:
+    """ALL_OF attribute with one pure ML-KEM-1024 key assigned at value level."""
+    return _create_keyed_attribute(
+        otdfctl,
+        temporary_namespace,
+        "mlkem1024-test",
+        [("m2", key_mlkem_1024)],
+        otdf_client_scs,
+        "mechanism-mlkem",
+    )
+
+
 # ---------------------------------------------------------------------------
 # Attribute + key assignment fixture (attribute-level)
 # ---------------------------------------------------------------------------

xtest/sdk/go/cli.sh

@@ -110,6 +110,11 @@ if [ "$1" == "supports" ]; then
       "${cmd[@]}" help policy kas-registry key create | grep -i hpqt:secp256r1-mlkem768
       exit $?
       ;;
+    mechanism-mlkem)
+      set -o pipefail
+      "${cmd[@]}" help policy kas-registry key create | grep -iE 'mlkem:768|mlkem:1024'
+      exit $?
+      ;;
     *)
       echo "Unknown feature: $2"
       exit 2

xtest/tdfs.py

@@ -60,6 +60,8 @@ def is_sdk_type(val: str) -> TypeIs[sdk_type]:
     "mechanism-xwing",
     # Support for encrypting with hybrid post-quantum/traditional KEM with NIST Elliptic Curves.
     "mechanism-secpmlkem",
+    # Support for pure (non-hybrid) ML-KEM key wrapping: mlkem:768 and mlkem:1024.
+    "mechanism-mlkem",
     "ns_grants",
     "obligations",
 ]
@@ -143,6 +145,11 @@ def __init__(self, **kwargs: dict[str, Any]):
             self.features.add("mechanism-xwing")
             self.features.add("mechanism-secpmlkem")
 
+        # Pure ML-KEM (non-hybrid) added by platform PR #3537. Tentatively v0.17.0;
+        # bump when the release target is finalized.
+        if self.semver >= (0, 17, 0):
+            self.features.add("mechanism-mlkem")
+
         print(f"PLATFORM_VERSION '{v}' supports [{', '.join(self.features)}]")
 
     def skip_if_unsupported(self, *features: feature_type):

xtest/test_pqc.py

@@ -19,6 +19,12 @@
 XWING_ENCAPSULATION_KEY_SIZE = 1216  # public key, bytes
 XWING_CIPHERTEXT_SIZE = 1120  # KEM ciphertext (wrappedKey), bytes
 
+# Pure ML-KEM sizes per FIPS 203 §8.
+MLKEM768_ENCAPSULATION_KEY_SIZE = 1184
+MLKEM768_CIPHERTEXT_SIZE = 1088
+MLKEM1024_ENCAPSULATION_KEY_SIZE = 1568
+MLKEM1024_CIPHERTEXT_SIZE = 1568
+
 
 def _b64_decoded_len(s: str) -> int:
     """Return the byte length of a base64-encoded string."""
@@ -231,6 +237,26 @@ def test_secpmlkem_3_roundtrip(
     assert filecmp.cmp(pt_file, rt_file)
 
 
+def _assert_mlkem_kao(
+    kao: KeyAccessObject,
+    expected_kids: set[str],
+    expected_url: str,
+    min_ciphertext_size: int,
+) -> None:
+    assert kao.kid in expected_kids
+    assert kao.url == expected_url
+    # PR #3537 currently emits the legacy "wrapped" type for pure ML-KEM. Some
+    # roadmap notes mention "mlkem-wrapped"; accept either so the test records
+    # which one ships without breaking on a naming change.
+    assert kao.type in {"wrapped", "mlkem-wrapped"}, (
+        f"unexpected KAO type for ML-KEM: {kao.type!r}"
+    )
+    wrapped_len = _b64_decoded_len(kao.wrappedKey)
+    assert wrapped_len > min_ciphertext_size, (
+        f"wrappedKey should exceed raw ML-KEM ciphertext ({min_ciphertext_size}), got {wrapped_len}"
+    )
+
+
 def test_secpmlkem_5_roundtrip(
     attribute_with_secpmlkem_5_key: tuple[Attribute, list[str]],
     key_secpmlkem_5: KasKey,
@@ -287,3 +313,105 @@ def test_secpmlkem_5_roundtrip(
     rt_file = encrypted_tdf.rt_file(ct_file, decrypt_sdk)
     decrypt_sdk.decrypt(ct_file, rt_file, "ztdf")
     assert filecmp.cmp(pt_file, rt_file)
+
+
+def test_mlkem_768_roundtrip(
+    attribute_with_mlkem_768_key: tuple[Attribute, list[str]],
+    key_mlkem_768: KasKey,
+    encrypt_sdk: tdfs.SDK,
+    decrypt_sdk: tdfs.SDK,
+    pt_file: Path,
+    kas_url_km1: str,
+    in_focus: set[tdfs.SDK],
+    encrypted_tdf: EncryptFactory,
+):
+    """Encrypt with a pure ML-KEM-768 managed key, then attempt decrypt.
+
+    The decrypt SDK is intentionally NOT pre-skipped on `mechanism-mlkem` so
+    we can observe whether SDKs lacking explicit pure-mlkem support can still
+    process the new ``mlkem-wrapped`` KAO type.
+    """
+    if not in_focus & {encrypt_sdk, decrypt_sdk}:
+        pytest.skip("Not in focus")
+    pfs = tdfs.get_platform_features()
+    pfs.skip_if_unsupported("key_management", "autoconfigure", "mechanism-mlkem")
+    encrypt_sdk.skip_if_unsupported(
+        "key_management", "autoconfigure", "mechanism-mlkem"
+    )
+    tdfs.skip_connectrpc_skew(encrypt_sdk, decrypt_sdk, pfs)
+    tdfs.skip_hexless_skew(encrypt_sdk, decrypt_sdk)
+
+    attr, key_ids = attribute_with_mlkem_768_key
+
+    ct_file = encrypted_tdf(
+        encrypt_sdk,
+        attr_values=attr.value_fqns,
+        target_mode=tdfs.select_target_version(encrypt_sdk, decrypt_sdk),
+    )
+
+    manifest = tdfs.manifest(ct_file)
+    assert len(manifest.encryptionInformation.keyAccess) == 1
+    _assert_mlkem_kao(
+        manifest.encryptionInformation.keyAccess[0],
+        expected_kids=set(key_ids),
+        expected_url=kas_url_km1,
+        min_ciphertext_size=MLKEM768_CIPHERTEXT_SIZE,
+    )
+    der_len = _pem_decoded_len(key_mlkem_768.key.public_key_ctx.pem)
+    assert der_len >= MLKEM768_ENCAPSULATION_KEY_SIZE, (
+        f"public key DER should be >= {MLKEM768_ENCAPSULATION_KEY_SIZE} bytes, got {der_len}"
+    )
+
+    rt_file = encrypted_tdf.rt_file(ct_file, decrypt_sdk)
+    decrypt_sdk.decrypt(ct_file, rt_file, "ztdf")
+    assert filecmp.cmp(pt_file, rt_file)
+
+
+def test_mlkem_1024_roundtrip(
+    attribute_with_mlkem_1024_key: tuple[Attribute, list[str]],
+    key_mlkem_1024: KasKey,
+    encrypt_sdk: tdfs.SDK,
+    decrypt_sdk: tdfs.SDK,
+    pt_file: Path,
+    kas_url_km1: str,
+    in_focus: set[tdfs.SDK],
+    encrypted_tdf: EncryptFactory,
+):
+    """Encrypt with a pure ML-KEM-1024 managed key, then attempt decrypt.
+
+    See ``test_mlkem_768_roundtrip`` for the decrypt-SDK skip rationale.
+    """
+    if not in_focus & {encrypt_sdk, decrypt_sdk}:
+        pytest.skip("Not in focus")
+    pfs = tdfs.get_platform_features()
+    pfs.skip_if_unsupported("key_management", "autoconfigure", "mechanism-mlkem")
+    encrypt_sdk.skip_if_unsupported(
+        "key_management", "autoconfigure", "mechanism-mlkem"
+    )
+    tdfs.skip_connectrpc_skew(encrypt_sdk, decrypt_sdk, pfs)
+    tdfs.skip_hexless_skew(encrypt_sdk, decrypt_sdk)
+
+    attr, key_ids = attribute_with_mlkem_1024_key
+
+    ct_file = encrypted_tdf(
+        encrypt_sdk,
+        attr_values=attr.value_fqns,
+        target_mode=tdfs.select_target_version(encrypt_sdk, decrypt_sdk),
+    )
+
+    manifest = tdfs.manifest(ct_file)
+    assert len(manifest.encryptionInformation.keyAccess) == 1
+    _assert_mlkem_kao(
+        manifest.encryptionInformation.keyAccess[0],
+        expected_kids=set(key_ids),
+        expected_url=kas_url_km1,
+        min_ciphertext_size=MLKEM1024_CIPHERTEXT_SIZE,
+    )
+    der_len = _pem_decoded_len(key_mlkem_1024.key.public_key_ctx.pem)
+    assert der_len >= MLKEM1024_ENCAPSULATION_KEY_SIZE, (
+        f"public key DER should be >= {MLKEM1024_ENCAPSULATION_KEY_SIZE} bytes, got {der_len}"
+    )
+
+    rt_file = encrypted_tdf.rt_file(ct_file, decrypt_sdk)
+    decrypt_sdk.decrypt(ct_file, rt_file, "ztdf")
+    assert filecmp.cmp(pt_file, rt_file)