test_n_generate_test_vectors.py

# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
"""Generate test vectors for use in testing the decrypt oracle."""
import base64
import binascii
import json
import os
from typing import Dict, Iterable, Text

import aws_encryption_sdk
import pytest
from aws_encryption_sdk.key_providers.base import MasterKeyProvider
from aws_encryption_sdk.key_providers.kms import KMSMasterKey
from aws_encryption_sdk_decrypt_oracle.key_providers.counting import CountingMasterKey
from aws_encryption_sdk_decrypt_oracle.key_providers.null import NullMasterKey

from .integration.integration_test_utils import test_vectors_filename

HERE = os.path.abspath(os.path.dirname(__file__))
GENERATE_VECTORS = "AWS_ENCRYPTION_SDK_PYTHON_DECRYPT_ORACLE_GENERATE_TEST_VECTORS"
PUBLIC_CMK = "arn:aws:kms:us-west-2:658956600833:alias/EncryptDecrypt"
ENCRYPTION_CONTEXT = {"key1": "val1", "key2": "val2"}


def _key_providers() -> Iterable[MasterKeyProvider]:
    """Generate all master key providers for test vector generation.
    Each will be used independently.
    """
    yield NullMasterKey()
    yield CountingMasterKey()
    yield KMSMasterKey(key_id=PUBLIC_CMK)


def _generate_vectors(key_provider: MasterKeyProvider, plaintext: bytes) -> Iterable[Dict[Text, Text]]:
    """Generate all desired test vectors for a given key provider and plaintext."""
    for algorithm_suite in aws_encryption_sdk.Algorithm:
        ciphertext, _header = aws_encryption_sdk.encrypt(
            source=plaintext,
            encryption_context=ENCRYPTION_CONTEXT,
            key_provider=key_provider,
            algorithm=algorithm_suite,
        )
        yield {
            "key-type": key_provider.provider_id,
            "algorithm-suite": binascii.hexlify(algorithm_suite.id_as_bytes()).decode("utf-8"),
            "ciphertext": base64.b64encode(ciphertext).decode("utf-8"),
            "plaintext": base64.b64encode(plaintext).decode("utf-8"),
        }


@pytest.mark.generate
@pytest.mark.skipif(GENERATE_VECTORS not in os.environ, reason="Generating test vectors is a rare occurance.")
def test_not_a_test_generate_test_vectors():
    """Generate all expected test vectors and write them to ``test/vectors/decrypt_oracle.json``."""
    vectors = []
    plaintext = os.urandom(64)
    for key_provider in _key_providers():
        vectors.extend(_generate_vectors(key_provider, plaintext))

    with open(test_vectors_filename(), "w") as output:
        json.dump(vectors, output, indent=4)