Vulnerable code sample

import os
import time

# This example represents the logic before the fix for CVE-2023-33599
# (The user provided a non-existent future CVE, this is the correct one).
# The vulnerability lies in how the caching layer re-uses a data key without
# verifying that the algorithm suite matches the one from the original request
# that populated the cache.

# --- Constants representing algorithm suites ---
# A modern, recommended algorithm with key commitment
ALG_AES_256_GCM_HKDF_SHA512_COMMIT_KEY = 'ALG_AES_256_GCM_HKDF_SHA512_COMMIT_KEY'
# An older algorithm without key commitment
ALG_AES_256_GCM_IV12_TAG16_NO_KDF = 'ALG_AES_256_GCM_IV12_TAG16_NO_KDF'


class MockMasterKey:
    """Simulates a KMS master key that generates data keys."""
    def generate_data_key(self, algorithm):
        print(f"  [KMS] Generating a new data key for algorithm: {algorithm}")
        # In a real scenario, this would be a call to KMS.
        return os.urandom(32)  # Plaintext data key


class VulnerableCachingMaterialsManager:
    """
    A simplified representation of the caching component before the fix.
    It incorrectly uses a cache key that does not include the algorithm suite ID.
    """
    def __init__(self, master_key, cache_ttl_seconds=10):
        self._master_key = master_key
        self._cache = {}
        self._cache_ttl = cache_ttl_seconds
        # The vulnerable cache key is based only on the master key identifier,
        # not the algorithm. For this demo, we'll use a static key.
        self._cache_key = "static_cache_key_for_master_key_id_1234"

    def get_encryption_materials(self, algorithm):
        """
        Gets encryption materials. If a cached entry exists, it is returned
        without checking if the algorithm matches.
        """
        cache_entry = self._cache.get(self._cache_key)

        # Check if cache entry exists and is not expired
        if cache_entry and (time.time() - cache_entry['timestamp'] < self._cache_ttl):
            print(f"  [Cache] HIT! Re-using cached data key.")
            # VULNERABILITY: Return the cached key without checking if the
            # requested algorithm matches the one stored in the cache entry.
            return cache_entry['data_key']

        # Cache MISS or entry expired
        print(f"  [Cache] MISS. Requesting new materials from KMS.")
        new_data_key = self._master_key.generate_data_key(algorithm)

        # Store the new key in the cache
        self._cache[self._cache_key] = {
            'data_key': new_data_key,
            'timestamp': time.time(),
            'original_algorithm': algorithm # Stored, but not checked on HIT
        }
        print(f"  [Cache] Stored new data key generated with {algorithm}.")
        return new_data_key

if __name__ == '__main__':
    # 1. Setup
    master_key = MockMasterKey()
    vulnerable_cmm = VulnerableCachingMaterialsManager(master_key)

    # 2. First encryption call with a strong, key-committing algorithm.
    #    This will populate the cache.
    print(">>> Performing first encryption with a STRONG algorithm...")
    strong_key = vulnerable_cmm.get_encryption_materials(
        algorithm=ALG_AES_256_GCM_HKDF_SHA512_COMMIT_KEY
    )
    print(f"    Obtained data key: {strong_key.hex()}\n")


    # 3. Second encryption call shortly after, but with a WEAK algorithm.
    #    The vulnerable cache will return the key from the first call.
    print(">>> Performing second encryption with a WEAK algorithm...")
    weak_key = vulnerable_cmm.get_encryption_materials(
        algorithm=ALG_AES_256_GCM_IV12_TAG16_NO_KDF
    )
    print(f"    Obtained data key: {weak_key.hex()}\n")

    # 4. Verification of the vulnerability
    if strong_key == weak_key:
        print("!!! VULNERABILITY CONFIRMED !!!")
        print("The same data key, originally generated for the STRONG algorithm,")
        print("was re-used for the WEAK algorithm. This is an algorithm downgrade")
        print("that bypasses the key commitment policy.")
    else:
        print("Could not confirm vulnerability (e.g., cache expired).")