chore(xtest): Add test of key mgmt create WIP

.github/workflows/xtest.yml

@@ -234,6 +234,38 @@ jobs:
           ec-tdf-enabled: true
           extra-keys: ${{ steps.load-extra-keys.outputs.EXTRA_KEYS }}
 
+      - name: Enable key management
+        run: |-
+          OT_CONFIG_FILE="$(pwd)/opentdf.yaml"
+          echo "OT_CONFIG_FILE=$OT_CONFIG_FILE">> "$GITHUB_ENV"
+          key_management_enabled=$(yq e '.services.kas.preview.key_management' "$OT_CONFIG_FILE")
+          case "$key_management_enabled" in
+            true)
+              echo "Key management is already enabled."
+              echo "OT_KEY_MANAGEMENT_ENABLED=true" >> "$GITHUB_ENV"
+              ;;
+            false)
+              echo "Enabling key management..."
+              yq eval -i '.services.kas.preview.key_management = true' opentdf.yaml
+              echo "OT_KEY_MANAGEMENT_ENABLED=true" >> "$GITHUB_ENV"
+              ;;
+            *)
+              echo "Key management is not available"
+              echo "OT_KEY_MANAGEMENT_ENABLED=false" >> "$GITHUB_ENV"
+              exit 0
+              ;;
+          esac
+          # Extract or add a new root key
+          root_key=$(yq e '.services.kas.root_key' "$OT_CONFIG_FILE")
+          if [ -z "$root_key" ]; then
+            echo "No root key found, generating a new one..."
+            root_key=$(openssl rand -hex 32)
+            yq eval -i ".services.kas.root_key = \"$root_key\"" "$OT_CONFIG_FILE"
+          fi
+          root_key=$(yq e '.services.kas.root_key' "$OT_CONFIG_FILE")
+          echo "OT_ROOT_KEY=$root_key" >> "$GITHUB_ENV"
+        working-directory: ${{ steps.run-platform.outputs.platform-working-dir }}
+
       - name: Set up Python 3.12
         uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b
         with:

xtest/README.md

@@ -60,6 +60,7 @@ go mod edit -replace github.com/opentdf/platform/lib/flattening=$GH_ORG_DIR/plat
 go mod edit -replace github.com/opentdf/platform/lib/ocrypto=$GH_ORG_DIR/platform/lib/ocrypto
 go mod edit -replace github.com/opentdf/platform/protocol/go=$GH_ORG_DIR/platform/protocol/go
 go mod edit -replace github.com/opentdf/platform/sdk=$GH_ORG_DIR/platform/sdk
+go mod tidy
 ```
 
 #### Build the SDKs

xtest/abac.py

@@ -2,6 +2,8 @@
 import json
 import logging
 import os
+import random
+import string
 import subprocess
 import sys
 import base64
@@ -46,12 +48,25 @@ class AttributeRule(enum.IntEnum):
     HIERARCHY = 3
 
 
+class SimpleKasPublicKey(BaseModelIgnoreExtra):
+    algorithm: int
+    kid: str
+    pem: str
+
+
+class SimpleKasKey(BaseModelIgnoreExtra):
+    kas_uri: str
+    public_key: SimpleKasPublicKey
+    kas_id: str
+
+
 class AttributeValue(BaseModelIgnoreExtra):
     id: str
     value: str
     fqn: str | None = None
     active: BoolValue | None = None
     metadata: Metadata | None = None
+    kas_keys: list[SimpleKasKey] | None = None
 
 
 class Attribute(BaseModelIgnoreExtra):
@@ -206,25 +221,39 @@ class KasPublicKey(BaseModelIgnoreExtra):
     algStr: str | None = Field(default=None, exclude=True)
 
 
+class PrivateKeyCtx(BaseModelIgnoreExtra):
+    key_id: str
+    wrapped_key: str
+
+
+class KeyProviderConfig(BaseModelIgnoreExtra):
+    id: str
+    name: str
+    config_json: bytes | None = None
+    metadata: Metadata | None = None
+
+
 # Helper model for the structure within key.public_key_ctx in the KAS key creation response
-class KasKeyResponsePublicKeyContext(BaseModelIgnoreExtra):
+class PublicKeyCtx(BaseModelIgnoreExtra):
     pem: str
 
 
 # Helper model for the nested "key" object in the KAS key creation response
-class KasKeyResponseKeyDetails(BaseModelIgnoreExtra):
+class AsymmetricKey(BaseModelIgnoreExtra):
     id: str
     key_id: str
     key_algorithm: int
     key_status: int
     key_mode: int
-    public_key_ctx: KasKeyResponsePublicKeyContext
+    public_key_ctx: PublicKeyCtx
+    private_key_ctx: PrivateKeyCtx | None = None
+    provider_config: KeyProviderConfig | None = None
     metadata: Metadata | None = None
 
 
 class KasKey(BaseModelIgnoreExtra):
     kas_id: str
-    key: KasKeyResponseKeyDetails
+    key: AsymmetricKey
     kas_uri: str
 
 
@@ -304,6 +333,32 @@ def kas_registry_create_if_not_present(
                 return e
         return self.kas_registry_create(uri, key)
 
+    def kas_registry_key_create(self, kas: KasEntry, key_id: str | None) -> KasKey:
+        cmd = self.otdfctl + "policy kas-registry key create".split()
+        if not key_id:
+            key_id = "".join(random.choices(string.ascii_lowercase, k=8))
+        cmd += [
+            f"--kas={kas.id}",
+            f"--key-id={key_id}",
+            "--mode=local",
+            "--algorithm=rsa:2048",
+            "--wrapping-key-id=wrapping-key-1",
+            f"--wrapping-key={os.environ['OT_ROOT_KEY']}",
+        ]
+
+        logger.info(f"kr-keys-create [{' '.join(cmd)}]")
+
+        cmd += [f"--wrapping-key={os.environ['OT_ROOT_KEY']}"]
+        process = subprocess.Popen(cmd, stdout=subprocess.PIPE)
+        out, err = process.communicate()
+        if err:
+            print(err, file=sys.stderr)
+            raise RuntimeError(f"Error creating KAS key: {err.decode()}")
+        if out:
+            print(out)
+        assert process.returncode == 0
+        return KasKey.model_validate_json(out)
+
     def kas_registry_keys_list(self, kas: KasEntry) -> list[KasKey]:
         cmd = self.otdfctl + "policy kas-registry key list".split()
         cmd += [f"--kas={kas.uri}"]

xtest/conftest.py

@@ -281,7 +281,7 @@ def kas_public_key_e1() -> abac.KasPublicKey:
 
 @pytest.fixture(scope="session")
 def kas_url_default():
-    return os.getenv("KASURL", "http://localhost:8080/kas")
+    return os.getenv("KASURL", "http://localhost:8080")
 
 
 @pytest.fixture(scope="module")
@@ -295,7 +295,7 @@ def kas_entry_default(
 
 @pytest.fixture(scope="session")
 def kas_url_value1():
-    return os.getenv("KASURL1", "http://localhost:8181/kas")
+    return os.getenv("KASURL1", "http://localhost:8181")
 
 
 @pytest.fixture(scope="module")
@@ -309,7 +309,7 @@ def kas_entry_value1(
 
 @pytest.fixture(scope="session")
 def kas_url_value2():
-    return os.getenv("KASURL2", "http://localhost:8282/kas")
+    return os.getenv("KASURL2", "http://localhost:8282")
 
 
 @pytest.fixture(scope="module")
@@ -323,7 +323,7 @@ def kas_entry_value2(
 
 @pytest.fixture(scope="session")
 def kas_url_attr():
-    return os.getenv("KASURL3", "http://localhost:8383/kas")
+    return os.getenv("KASURL3", "http://localhost:8383")
 
 
 @pytest.fixture(scope="module")
@@ -337,7 +337,7 @@ def kas_entry_attr(
 
 @pytest.fixture(scope="session")
 def kas_url_ns():
-    return os.getenv("KASURL4", "http://localhost:8484/kas")
+    return os.getenv("KASURL4", "http://localhost:8484")
 
 
 @pytest.fixture(scope="module")
@@ -422,6 +422,54 @@ def attribute_with_different_kids(
     return allof
 
 
+@pytest.fixture(scope="module")
+def attribute_with_managed_keys(
+    otdfctl: abac.OpentdfCommandLineTool,
+    kas_entry_default: abac.KasEntry,
+    temporary_namespace: abac.Namespace,
+    otdf_client_scs: abac.SubjectConditionSet,
+):
+    """
+    Create an attribute with a newly created managed key.
+    """
+    pfs = tdfs.PlatformFeatureSet()
+    if "key_management" not in pfs.features:
+        pytest.skip(
+            "Key management feature is not enabled, skipping test for multiple KAS keys"
+        )
+
+    managed_key = otdfctl.kas_registry_key_create(kas_entry_default, None)
+
+    allof = otdfctl.attribute_create(
+        temporary_namespace,
+        "managedkeys",
+        abac.AttributeRule.ALL_OF,
+        [managed_key.key.key_id],
+    )
+    assert allof.values
+    (ar1,) = allof.values
+    assert ar1.value == managed_key.key.key_id
+
+    sm = otdfctl.scs_map(otdf_client_scs, ar1)
+    assert sm.attribute_value.value == ar1.value
+
+    # Assign kas key to the attribute values
+    otdfctl.key_assign_value(managed_key, ar1)
+    ar1.kas_keys = [
+        abac.SimpleKasKey(
+            kas_uri=managed_key.kas_uri,
+            public_key=abac.SimpleKasPublicKey(
+                algorithm=managed_key.key.key_algorithm,
+                kid=managed_key.key.key_id,
+                pem=managed_key.key.public_key_ctx.pem,
+            ),
+            kas_id=managed_key.kas_id,
+        )
+    ]
+
+    return allof
+
+
 @pytest.fixture(scope="module")
 def attribute_single_kas_grant(
     otdfctl: abac.OpentdfCommandLineTool,

xtest/test_abac.py

@@ -71,7 +71,60 @@ def test_key_mapping_multiple_mechanisms(
     assert manifest.encryptionInformation.keyAccess[0].url == kas_url_default
 
     tdfs.skip_if_unsupported(decrypt_sdk, "ecwrap")
-    rt_file = tmp_dir / f"multimechanism-{encrypt_sdk}-{decrypt_sdk}.untdf"
+    rt_file = tmp_dir / f"{sample_name}-{decrypt_sdk}.untdf"
+    decrypt_sdk.decrypt(ct_file, rt_file, "ztdf")
+    assert filecmp.cmp(pt_file, rt_file)
+
+
+def test_key_mapping_from_mgmt(
+    attribute_with_managed_keys: Attribute,
+    encrypt_sdk: tdfs.SDK,
+    decrypt_sdk: tdfs.SDK,
+    tmp_dir: Path,
+    pt_file: Path,
+    kas_url_default: str,
+    in_focus: set[tdfs.SDK],
+):
+    global counter
+
+    tdfs.skip_if_unsupported(encrypt_sdk, "key_management")
+    skip_dspx1153(encrypt_sdk, decrypt_sdk)
+    if not in_focus & {encrypt_sdk, decrypt_sdk}:
+        pytest.skip("Not in focus")
+    tdfs.skip_if_unsupported(encrypt_sdk, "autoconfigure")
+    pfs = tdfs.PlatformFeatureSet()
+    tdfs.skip_connectrpc_skew(encrypt_sdk, decrypt_sdk, pfs)
+    tdfs.skip_hexless_skew(encrypt_sdk, decrypt_sdk)
+
+    sample_name = f"from-mgmt-{encrypt_sdk}"
+    if sample_name in cipherTexts:
+        ct_file = cipherTexts[sample_name]
+    else:
+        ct_file = tmp_dir / f"{sample_name}.tdf"
+        cipherTexts[sample_name] = ct_file
+        # Currently, we only support rsa:2048 and ec:secp256r1
+        encrypt_sdk.encrypt(
+            pt_file,
+            ct_file,
+            mime_type="text/plain",
+            container="ztdf",
+            attr_values=attribute_with_managed_keys.value_fqns,
+            target_mode=tdfs.select_target_version(encrypt_sdk, decrypt_sdk),
+        )
+
+    assert attribute_with_managed_keys.values
+    val = attribute_with_managed_keys.values[0]
+    assert val.kas_keys
+    assert len(val.kas_keys) == 1
+    kek = val.kas_keys[0]
+    manifest = tdfs.manifest(ct_file)
+    assert set([kao.kid for kao in manifest.encryptionInformation.keyAccess]) == set(
+        [kek.public_key.kid]
+    )
+    assert manifest.encryptionInformation.keyAccess[0].url == kas_url_default
+
+    tdfs.skip_if_unsupported(decrypt_sdk, "ecwrap")
+    rt_file = tmp_dir / f"{sample_name}-{decrypt_sdk}.untdf"
     decrypt_sdk.decrypt(ct_file, rt_file, "ztdf")
     assert filecmp.cmp(pt_file, rt_file)
 

xtest/test_nano.py

@@ -18,12 +18,12 @@ def test_magic_version():
 
 
 def test_resource_locator():
-    rl0 = nano.locator("https://localhost:8080/kas")
+    rl0 = nano.locator("https://localhost:8080")
     print(rl0)
     expected_bits = "01 12 6c 6f 63 61 6c 68 6f 73 74 3a 38 30 38 30 2f 6b 61 73"
     assert expected_bits == enc_hex(bytes(rl0))
 
-    rl1 = nano.locator("https://localhost:8080/kas", b"ab")
+    rl1 = nano.locator("https://localhost:8080", b"ab")
     print(rl1)
     assert """11 12 6c 6f 63 61 6c 68 6f 73 74 3a 38 30 38 30 2f 6b 61 73
 61 62""" == enc_hex(

xtest/test_tdfs.py

@@ -499,7 +499,7 @@ def change_payload_end(payload_bytes: bytes) -> bytes:
 def malicious_kao(manifest: tdfs.Manifest) -> tdfs.Manifest:
     assert manifest.encryptionInformation.keyAccess
     manifest.encryptionInformation.keyAccess[0].url = (
-        "http://localhost:8585/malicious/kas"  # nothing running at 8585
+        "http://localhost:8585/malicious"  # nothing running at 8585
     )
     return manifest