Extend tests for encryption commands.

Ref: NCSDK-31673

Signed-off-by: Artur Hadasz <[email protected]>

Author: ahasztag

ncs/encrypt_script.py

@@ -96,7 +96,6 @@ def generate_kms_artifacts(self, asset_plaintext: bytes, key_name: str, context:
 
         This method encrypts asset_plaintext bytes using the specified key wrap algorithm,
         and returns the encrypted asset and encrypted content encryption key (CEK).
-
         """
         # Enc structure:
         # {
@@ -255,7 +254,7 @@ def generate(
         if kw_alg == SuitKWAlgorithms.A256KW:
             if encrypted_cek is None:
                 raise ValueError("Encrypted CEK is required for AES Key Wrap 256")
-        self.kw_alg_convert(kw_alg)
+        self._kw_alg_convert(kw_alg)
         return self.generate_encryption_info_and_encrypted_payload(encrypted_asset, encrypted_cek, key_id)
 
 

suit_generator/cmd_encrypt.py

@@ -175,8 +175,8 @@ def encrypt_and_generate(**kwargs):
         kwargs["key_name"],
         kwargs["key_id"],
         kwargs["context"],
-        kwargs["hash_alg"],
-        kwargs["kw_alg"],
+        SuitDigestAlgorithms(kwargs["hash_alg"]),
+        SuitKWAlgorithms(kwargs["kw_alg"]),
         kwargs["kms_script"],
     )
     with open(os.path.join(kwargs["output_dir"], "plain_text_digest.bin"), "wb") as file:
@@ -200,7 +200,7 @@ def generate_info(**kwargs):
         encrypted_firmware,
         encrypted_key,
         kwargs["key_id"],
-        kwargs["kw_alg"],
+        SuitKWAlgorithms(kwargs["kw_alg"]),
     )
     with open(os.path.join(kwargs["output_dir"], "suit_encryption_info.bin"), "wb") as file:
         file.write(encryption_info)

tests/encrypt_script_mock.py

@@ -0,0 +1,78 @@
+#
+# Copyright (c) 2025 Nordic Semiconductor ASA
+#
+# SPDX-License-Identifier: LicenseRef-Nordic-5-Clause
+#
+"""Script mocking the Encrypt script."""
+import json
+
+from pathlib import Path
+from suit_generator.suit_encrypt_script_base import (
+    SuitEncryptorBase,
+    SuitDigestAlgorithms,
+    SuitKWAlgorithms,
+)
+
+
+class EncryptorMock(SuitEncryptorBase):
+    """Encryptor mock implementation."""
+
+    def encrypt_and_generate(
+        self,
+        firmware: bytes,
+        key_name: str,
+        key_id: int,
+        context: str,
+        hash_alg: SuitDigestAlgorithms,
+        kw_alg: SuitKWAlgorithms,
+        kms_script: Path,
+    ) -> tuple[bytes, bytes, bytes, bytes, int]:
+        """Mock encrypting the payload and generation of encryption artifacts."""
+        context_loaded = json.loads(context)
+        self.output_file = f"test_output_{key_name}.json"
+        self.json_data = {"firmware": firmware.hex()}
+        self.json_data["key_name"] = key_name
+        self.json_data["key_id"] = key_id
+        self.json_data["kw_alg"] = kw_alg.value
+        self.json_data["hash_alg"] = hash_alg.value
+        self.json_data["context"] = context_loaded["ctx"]
+        self.json_data["kms_script"] = kms_script
+        with open(self.output_file, "w") as f:
+            json.dump(self.json_data, f)
+
+        return (
+            bytes.fromhex(context_loaded["encrypted_data"]),
+            bytes.fromhex(context_loaded["tag"]),
+            bytes.fromhex(context_loaded["encryption_info"]),
+            bytes.fromhex(context_loaded["digest"]),
+            context_loaded["plaintext_length"],
+        )
+
+    def generate(
+        self, encrypted_asset: bytes, encrypted_cek: bytes, key_id: int, kw_alg: SuitKWAlgorithms
+    ) -> tuple[bytes, bytes, bytes]:
+        """
+        Mock generation of encryption artifacts.
+
+        The encrypted asset for this mock is a json object, containing the return values.
+        """
+        #         :return: The encrypted payload, tag, encryption info.
+        context_loaded = json.loads(encrypted_asset.decode())
+        self.output_file = f"test_output_{key_id}.json"
+
+        self.json_data = {"key_id": key_id}
+        self.json_data["kw_alg"] = kw_alg.value
+        self.json_data["encrypted_cek"] = encrypted_cek.hex()
+        with open(self.output_file, "w") as f:
+            json.dump(self.json_data, f)
+
+        return (
+            bytes.fromhex(context_loaded["encrypted_data"]),
+            bytes.fromhex(context_loaded["tag"]),
+            bytes.fromhex(context_loaded["encryption_info"]),
+        )
+
+
+def suit_encryptor_factory():
+    """Get an Encryptor object."""
+    return EncryptorMock()

tests/kms_script_mock.py

@@ -32,16 +32,17 @@ def init_kms(self, context: str) -> None:
     def encrypt(self, plaintext: bytes, key_name: str, context: str, aad: bytes) -> tuple[bytes, bytes, bytes]:
         """Mock of the KMS script encrypt function."""
         context_loaded = json.loads(context)
-        self.json_data["encrypt_plaintext"] = plaintext.decode()
+        self.json_data["encrypt_plaintext"] = plaintext.hex()
         self.json_data["encrypt_key_name"] = key_name
         self.json_data["encrypt_context"] = context_loaded["ctx"]
-        self.json_data["encrypt_aad"] = aad.decode()
+        self.json_data["encrypt_aad"] = aad.hex()
         with open(self.output_file, "w") as f:
             json.dump(self.json_data, f)
+
         return (
-            context_loaded["iv"].encode(),
-            context_loaded["encryption_key"].encode(),
-            context_loaded["encrypted_data"].encode(),
+            bytes.fromhex(context_loaded["iv"]),
+            bytes.fromhex(context_loaded["tag"]),
+            bytes.fromhex(context_loaded["encrypted_data"]),
         )
 
     def sign(self, data: bytes, key_name: str, algorithm: str, context: str) -> bytes:

tests/test_cmd_encrypt.py

@@ -0,0 +1,220 @@
+# Copyright (c) 2025 Nordic Semiconductor ASA
+#
+# SPDX-License-Identifier: LicenseRef-Nordic-5-Clause
+#
+"""Unit tests for cmd_encrypt.py implementation."""
+
+import pytest
+import os
+import pathlib
+import encrypt_script_mock
+import json
+
+from suit_generator.cmd_encrypt import main as cmd_encrypt_main
+
+
+TEMP_DIRECTORY = pathlib.Path("test_test_data")
+
+
+@pytest.fixture
+def setup_and_teardown(tmp_path_factory):
+    """Create and cleanup environment."""
+    # Setup environment
+    #   - create required files in TEMP_DIRECTORY
+    start_directory = os.getcwd()
+    path = tmp_path_factory.mktemp(TEMP_DIRECTORY)
+    os.chdir(path)
+    yield
+    # Cleanup environment
+    #   - remove temp directory
+    os.chdir(start_directory)
+
+
+@pytest.mark.parametrize(
+    "firmware, key_name, key_id, hash_alg, kw_alg, ctx, kms_script, encrypted_data, "
+    + "tag, encryption_info, digest, plaintext_length, output_dir_name",
+    [
+        (
+            b"test_firmware",
+            "test_key",
+            0x40000000,
+            "sha-256",
+            "direct",
+            "test_ctx",
+            "test_kms_script.py",
+            b"test_encrypted_data",
+            b"test_tag",
+            b"test_encryption_info",
+            b"test_digest",
+            123,
+            "enc_artifacts1",
+        ),
+        (
+            b"firmware2",
+            "key2",
+            0x12345678,
+            "sha-512",
+            "aes-kw-256",
+            "testctx2",
+            "kms_script2.py",
+            b"encrypted_data2",
+            b"tag2",
+            b"encryption_info2",
+            b"digest2",
+            5000,
+            "enc_artifacts2",
+        ),
+    ],
+)
+def test_encrypt_cmd_encrypt_and_generate(
+    setup_and_teardown,
+    firmware,
+    key_name,
+    key_id,
+    hash_alg,
+    kw_alg,
+    ctx,
+    kms_script,
+    encrypted_data,
+    tag,
+    encryption_info,
+    digest,
+    plaintext_length,
+    output_dir_name,
+):
+    """Test the encrypt-and-generate command"""
+
+    full_context = json.dumps(
+        {
+            "ctx": ctx,
+            "encrypted_data": encrypted_data.hex(),
+            "tag": tag.hex(),
+            "encryption_info": encryption_info.hex(),
+            "digest": digest.hex(),
+            "plaintext_length": plaintext_length,
+        }
+    )
+
+    os.mkdir(output_dir_name)
+    firmware_file_name = "fw.bin"
+    with open(firmware_file_name, "wb") as file:
+        file.write(firmware)
+
+    kwargs = {
+        "encrypt_subcommand": "encrypt-and-generate",
+        "firmware": firmware_file_name,
+        "key_name": key_name,
+        "key_id": key_id,
+        "hash_alg": hash_alg,
+        "kw_alg": kw_alg,
+        "context": full_context,
+        "output_dir": output_dir_name,
+        "kms_script": kms_script,
+        "encrypt_script": str(encrypt_script_mock.__file__),
+    }
+
+    cmd_encrypt_main(**kwargs)
+
+    assert os.path.exists(f"test_output_{key_name}.json")
+    assert os.path.exists(output_dir_name)
+    assert os.path.exists(f"{output_dir_name}/plain_text_digest.bin")
+    assert os.path.exists(f"{output_dir_name}/plain_text_size.txt")
+    assert os.path.exists(f"{output_dir_name}/suit_encryption_info.bin")
+    assert os.path.exists(f"{output_dir_name}/encrypted_content.bin")
+
+    test_json_data = json.load(open(f"test_output_{key_name}.json"))
+
+    assert test_json_data["firmware"] == firmware.hex()
+    assert test_json_data["key_name"] == key_name
+    assert test_json_data["key_id"] == key_id
+    assert test_json_data["kw_alg"] == kw_alg
+    assert test_json_data["hash_alg"] == hash_alg
+    assert test_json_data["context"] == ctx
+    assert test_json_data["kms_script"] == kms_script
+
+    assert open(f"{output_dir_name}/plain_text_digest.bin", "rb").read() == digest
+    assert open(f"{output_dir_name}/plain_text_size.txt", "r").read() == str(plaintext_length)
+    assert open(f"{output_dir_name}/suit_encryption_info.bin", "rb").read() == encryption_info
+    assert open(f"{output_dir_name}/encrypted_content.bin", "rb").read() == tag + encrypted_data
+
+
+@pytest.mark.parametrize(
+    "encrypted_firmware, encrypted_key, key_id, kw_alg, encrypted_data, tag, encryption_info, output_dir_name",
+    [
+        (
+            b"test_encrypted_firmware",
+            b"test_encrypted_key",
+            0x40000000,
+            "direct",
+            b"test_encrypted_data",
+            b"test_tag",
+            b"test_encryption_info",
+            "gen_artifacts1",
+        ),
+        (
+            b"encrypted_firmware2",
+            b"encrypted_key2",
+            0x12345678,
+            "aes-kw-256",
+            b"encrypted_data2",
+            b"tag2",
+            b"encryption_info2",
+            "gen_artifacts2",
+        ),
+    ],
+)
+def test_encrypt_cmd_generate(
+    setup_and_teardown,
+    encrypted_firmware,
+    encrypted_key,
+    key_id,
+    kw_alg,
+    encrypted_data,
+    tag,
+    encryption_info,
+    output_dir_name,
+):
+    """Test the generate-info command"""
+
+    os.mkdir(output_dir_name)
+    encrypted_firmware_file_name = "encrypted_fw.bin"
+    encrypted_key_file_name = "encrypted_key.bin"
+
+    context = json.dumps(
+        {
+            "encrypted_data": encrypted_data.hex(),
+            "tag": tag.hex(),
+            "encryption_info": encryption_info.hex(),
+        }
+    )
+
+    with open(encrypted_firmware_file_name, "wb") as file:
+        file.write(context.encode())
+    with open(encrypted_key_file_name, "wb") as file:
+        file.write(encrypted_key)
+
+    kwargs = {
+        "encrypt_subcommand": "generate-info",
+        "encrypted_firmware": encrypted_firmware_file_name,
+        "encrypted_key": encrypted_key_file_name,
+        "key_id": key_id,
+        "kw_alg": kw_alg,
+        "output_dir": output_dir_name,
+        "encrypt_script": str(encrypt_script_mock.__file__),
+    }
+
+    cmd_encrypt_main(**kwargs)
+
+    assert os.path.exists(f"test_output_{key_id}.json")
+    assert os.path.exists(output_dir_name)
+    assert os.path.exists(f"{output_dir_name}/suit_encryption_info.bin")
+    assert os.path.exists(f"{output_dir_name}/encrypted_content.bin")
+
+    test_json_data = json.load(open(f"test_output_{key_id}.json"))
+
+    assert test_json_data["key_id"] == key_id
+    assert test_json_data["kw_alg"] == kw_alg
+    assert test_json_data["encrypted_cek"] == encrypted_key.hex()
+
+    assert open(f"{output_dir_name}/suit_encryption_info.bin", "rb").read() == encryption_info
+    assert open(f"{output_dir_name}/encrypted_content.bin", "rb").read() == tag + encrypted_data