Cryptography and Key Management

/**
 * Synchronous key encryption key interface for cryptographic operations.
 */
interface KeyEncryptionKey {
    /**
     * Gets the identifier of the key encryption key.
     * @return The key identifier
     */
    String getKeyId();
    
    /**
     * Encrypts the specified plaintext.
     * @param algorithm The encryption algorithm to use
     * @param plaintext The plaintext to encrypt
     * @return The encrypted data (ciphertext)
     */
    byte[] wrapKey(String algorithm, byte[] plaintext);
    
    /**
     * Decrypts the specified ciphertext.
     * @param algorithm The decryption algorithm to use
     * @param ciphertext The ciphertext to decrypt
     * @return The decrypted data (plaintext)
     */
    byte[] unwrapKey(String algorithm, byte[] ciphertext);
}

/**
 * Asynchronous key encryption key interface for cryptographic operations.
 */
interface AsyncKeyEncryptionKey {
    /**
     * Gets the identifier of the key encryption key.
     * @return A Mono containing the key identifier
     */
    Mono<String> getKeyId();
    
    /**
     * Encrypts the specified plaintext asynchronously.
     * @param algorithm The encryption algorithm to use
     * @param plaintext The plaintext to encrypt
     * @return A Mono containing the encrypted data (ciphertext)
     */
    Mono<byte[]> wrapKey(String algorithm, byte[] plaintext);
    
    /**
     * Decrypts the specified ciphertext asynchronously.
     * @param algorithm The decryption algorithm to use  
     * @param ciphertext The ciphertext to decrypt
     * @return A Mono containing the decrypted data (plaintext)
     */
    Mono<byte[]> unwrapKey(String algorithm, byte[] ciphertext);
}

/**
 * Synchronous resolver interface for key encryption keys.
 */
interface KeyEncryptionKeyResolver {
    /**
     * Resolves a key encryption key by its identifier.
     * @param keyId The identifier of the key to resolve
     * @return The resolved KeyEncryptionKey, or null if not found
     */
    KeyEncryptionKey resolveKey(String keyId);
}

/**
 * Asynchronous resolver interface for key encryption keys.
 */
interface AsyncKeyEncryptionKeyResolver {
    /**
     * Resolves a key encryption key by its identifier asynchronously.
     * @param keyId The identifier of the key to resolve
     * @return A Mono containing the resolved AsyncKeyEncryptionKey, or empty if not found
     */
    Mono<AsyncKeyEncryptionKey> resolveKey(String keyId);
}

import com.azure.core.cryptography.*;
import javax.crypto.Cipher;
import javax.crypto.spec.SecretKeySpec;
import java.security.Key;

class AesKeyEncryptionKey implements KeyEncryptionKey {
    private final String keyId;
    private final Key key;
    
    public AesKeyEncryptionKey(String keyId, byte[] keyBytes) {
        this.keyId = keyId;
        this.key = new SecretKeySpec(keyBytes, "AES");
    }
    
    @Override
    public String getKeyId() {
        return keyId;
    }
    
    @Override
    public byte[] wrapKey(String algorithm, byte[] plaintext) {
        try {
            Cipher cipher = Cipher.getInstance(algorithm);
            cipher.init(Cipher.ENCRYPT_MODE, key);
            return cipher.doFinal(plaintext);
        } catch (Exception e) {
            throw new RuntimeException("Encryption failed", e);
        }
    }
    
    @Override
    public byte[] unwrapKey(String algorithm, byte[] ciphertext) {
        try {
            Cipher cipher = Cipher.getInstance(algorithm);
            cipher.init(Cipher.DECRYPT_MODE, key);
            return cipher.doFinal(ciphertext);
        } catch (Exception e) {
            throw new RuntimeException("Decryption failed", e);
        }
    }
}

import com.azure.core.cryptography.*;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

class AsyncAesKeyEncryptionKey implements AsyncKeyEncryptionKey {
    private final KeyEncryptionKey syncKey;
    
    public AsyncAesKeyEncryptionKey(KeyEncryptionKey syncKey) {
        this.syncKey = syncKey;
    }
    
    @Override
    public Mono<String> getKeyId() {
        return Mono.fromCallable(syncKey::getKeyId)
                   .subscribeOn(Schedulers.boundedElastic());
    }
    
    @Override
    public Mono<byte[]> wrapKey(String algorithm, byte[] plaintext) {
        return Mono.fromCallable(() -> syncKey.wrapKey(algorithm, plaintext))
                   .subscribeOn(Schedulers.boundedElastic());
    }
    
    @Override
    public Mono<byte[]> unwrapKey(String algorithm, byte[] ciphertext) {
        return Mono.fromCallable(() -> syncKey.unwrapKey(algorithm, ciphertext))
                   .subscribeOn(Schedulers.boundedElastic());
    }
}

import com.azure.core.cryptography.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.Map;

class InMemoryKeyResolver implements KeyEncryptionKeyResolver {
    private final Map<String, KeyEncryptionKey> keys = new ConcurrentHashMap<>();
    
    public void addKey(String keyId, KeyEncryptionKey key) {
        keys.put(keyId, key);
    }
    
    @Override
    public KeyEncryptionKey resolveKey(String keyId) {
        return keys.get(keyId);
    }
}

class AsyncInMemoryKeyResolver implements AsyncKeyEncryptionKeyResolver {
    private final KeyEncryptionKeyResolver syncResolver;
    
    public AsyncInMemoryKeyResolver(KeyEncryptionKeyResolver syncResolver) {
        this.syncResolver = syncResolver;
    }
    
    @Override
    public Mono<AsyncKeyEncryptionKey> resolveKey(String keyId) {
        return Mono.fromCallable(() -> {
            KeyEncryptionKey syncKey = syncResolver.resolveKey(keyId);
            return syncKey != null ? new AsyncAesKeyEncryptionKey(syncKey) : null;
        }).subscribeOn(Schedulers.boundedElastic());
    }
}

import com.azure.core.cryptography.*;
import com.azure.security.keyvault.keys.cryptography.*;
import com.azure.identity.DefaultAzureCredentialBuilder;

// Example using Azure Key Vault for key encryption
class KeyVaultKeyEncryptionKey implements KeyEncryptionKey {
    private final CryptographyClient cryptoClient;
    private final String keyId;
    
    public KeyVaultKeyEncryptionKey(String keyVaultUrl, String keyName) {
        CryptographyClientBuilder builder = new CryptographyClientBuilder()
            .keyIdentifier(keyVaultUrl + "/keys/" + keyName)
            .credential(new DefaultAzureCredentialBuilder().build());
        
        this.cryptoClient = builder.buildClient();
        this.keyId = keyName;
    }
    
    @Override
    public String getKeyId() {
        return keyId;
    }
    
    @Override
    public byte[] wrapKey(String algorithm, byte[] plaintext) {
        EncryptionAlgorithm encryptAlg = EncryptionAlgorithm.fromString(algorithm);
        EncryptResult result = cryptoClient.encrypt(encryptAlg, plaintext);
        return result.getCipherText();
    }
    
    @Override
    public byte[] unwrapKey(String algorithm, byte[] ciphertext) {
        EncryptionAlgorithm encryptAlg = EncryptionAlgorithm.fromString(algorithm);
        DecryptResult result = cryptoClient.decrypt(encryptAlg, ciphertext);
        return result.getPlainText();
    }
}

// Async version for Key Vault
class AsyncKeyVaultKeyEncryptionKey implements AsyncKeyEncryptionKey {
    private final CryptographyAsyncClient cryptoClient;
    private final String keyId;
    
    public AsyncKeyVaultKeyEncryptionKey(String keyVaultUrl, String keyName) {
        CryptographyClientBuilder builder = new CryptographyClientBuilder()
            .keyIdentifier(keyVaultUrl + "/keys/" + keyName)
            .credential(new DefaultAzureCredentialBuilder().build());
        
        this.cryptoClient = builder.buildAsyncClient();
        this.keyId = keyName;
    }
    
    @Override
    public Mono<String> getKeyId() {
        return Mono.just(keyId);
    }
    
    @Override
    public Mono<byte[]> wrapKey(String algorithm, byte[] plaintext) {
        EncryptionAlgorithm encryptAlg = EncryptionAlgorithm.fromString(algorithm);
        return cryptoClient.encrypt(encryptAlg, plaintext)
                          .map(EncryptResult::getCipherText);
    }
    
    @Override
    public Mono<byte[]> unwrapKey(String algorithm, byte[] ciphertext) {
        EncryptionAlgorithm encryptAlg = EncryptionAlgorithm.fromString(algorithm);
        return cryptoClient.decrypt(encryptAlg, ciphertext)
                          .map(DecryptResult::getPlainText);
    }
}

import com.azure.core.cryptography.*;
import com.azure.core.util.BinaryData;

class EncryptedDataService {
    private final AsyncKeyEncryptionKey keyEncryptionKey;
    private final AsyncKeyEncryptionKeyResolver keyResolver;
    
    public EncryptedDataService(AsyncKeyEncryptionKey kek, AsyncKeyEncryptionKeyResolver resolver) {
        this.keyEncryptionKey = kek;
        this.keyResolver = resolver;
    }
    
    public Mono<EncryptedBlob> encryptData(BinaryData data) {
        // Generate a content encryption key (CEK)
        byte[] contentKey = generateRandomKey(32); // 256-bit key
        
        return keyEncryptionKey.getKeyId()
            .flatMap(keyId -> keyEncryptionKey.wrapKey("AES/GCM/NoPadding", contentKey)
                .map(encryptedKey -> {
                    // Encrypt the actual data with the CEK
                    byte[] encryptedData = encryptWithContentKey(contentKey, data.toBytes());
                    
                    return new EncryptedBlob(keyId, encryptedKey, encryptedData);
                }));
    }
    
    public Mono<BinaryData> decryptData(EncryptedBlob encryptedBlob) {
        return keyResolver.resolveKey(encryptedBlob.getKeyId())
            .switchIfEmpty(Mono.error(new IllegalArgumentException("Key not found: " + encryptedBlob.getKeyId())))
            .flatMap(kek -> kek.unwrapKey("AES/GCM/NoPadding", encryptedBlob.getEncryptedKey()))
            .map(contentKey -> {
                // Decrypt the data with the recovered content key
                byte[] decryptedData = decryptWithContentKey(contentKey, encryptedBlob.getEncryptedData());
                return BinaryData.fromBytes(decryptedData);
            });
    }
    
    private byte[] generateRandomKey(int size) {
        byte[] key = new byte[size];
        new java.security.SecureRandom().nextBytes(key);
        return key;
    }
    
    private byte[] encryptWithContentKey(byte[] key, byte[] data) {
        // Implementation of AES encryption
        // This is a simplified example - real implementation would handle IV, etc.
        try {
            SecretKeySpec keySpec = new SecretKeySpec(key, "AES");
            Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
            cipher.init(Cipher.ENCRYPT_MODE, keySpec);
            return cipher.doFinal(data);
        } catch (Exception e) {
            throw new RuntimeException("Content encryption failed", e);
        }
    }
    
    private byte[] decryptWithContentKey(byte[] key, byte[] encryptedData) {
        // Implementation of AES decryption
        try {
            SecretKeySpec keySpec = new SecretKeySpec(key, "AES");
            Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
            cipher.init(Cipher.DECRYPT_MODE, keySpec);
            return cipher.doFinal(encryptedData);
        } catch (Exception e) {
            throw new RuntimeException("Content decryption failed", e);
        }
    }
}

// Helper class for encrypted data
class EncryptedBlob {
    private final String keyId;
    private final byte[] encryptedKey;
    private final byte[] encryptedData;
    
    public EncryptedBlob(String keyId, byte[] encryptedKey, byte[] encryptedData) {
        this.keyId = keyId;
        this.encryptedKey = encryptedKey.clone();
        this.encryptedData = encryptedData.clone();
    }
    
    public String getKeyId() { return keyId; }
    public byte[] getEncryptedKey() { return encryptedKey.clone(); }
    public byte[] getEncryptedData() { return encryptedData.clone(); }
}

import com.azure.core.cryptography.*;
import java.util.List;
import java.util.concurrent.ThreadLocalRandom;

class MultiKeyEncryptionService {
    private final List<AsyncKeyEncryptionKey> availableKeys;
    private final AsyncKeyEncryptionKeyResolver keyResolver;
    
    public MultiKeyEncryptionService(List<AsyncKeyEncryptionKey> keys, AsyncKeyEncryptionKeyResolver resolver) {
        this.availableKeys = keys;
        this.keyResolver = resolver;
    }
    
    public Mono<EncryptedBlob> encryptWithRandomKey(BinaryData data) {
        // Select a random key for encryption
        int keyIndex = ThreadLocalRandom.current().nextInt(availableKeys.size());
        AsyncKeyEncryptionKey selectedKey = availableKeys.get(keyIndex);
        
        return encryptWithSpecificKey(selectedKey, data);
    }
    
    public Mono<EncryptedBlob> encryptWithSpecificKey(AsyncKeyEncryptionKey kek, BinaryData data) {
        byte[] contentKey = generateRandomKey(32);
        
        return kek.getKeyId()
            .flatMap(keyId -> kek.wrapKey("RSA-OAEP", contentKey)
                .map(encryptedKey -> {
                    byte[] encryptedData = encryptWithContentKey(contentKey, data.toBytes());
                    return new EncryptedBlob(keyId, encryptedKey, encryptedData);
                }));
    }
    
    public Flux<BinaryData> decryptMultiple(List<EncryptedBlob> encryptedBlobs) {
        return Flux.fromIterable(encryptedBlobs)
            .flatMap(blob -> keyResolver.resolveKey(blob.getKeyId())
                .flatMap(kek -> kek.unwrapKey("RSA-OAEP", blob.getEncryptedKey()))
                .map(contentKey -> {
                    byte[] decryptedData = decryptWithContentKey(contentKey, blob.getEncryptedData());
                    return BinaryData.fromBytes(decryptedData);
                })
                .onErrorResume(error -> {
                    // Log error and continue with next blob
                    System.err.println("Failed to decrypt blob with key " + blob.getKeyId() + ": " + error.getMessage());
                    return Mono.empty();
                }));
    }
    
    private byte[] generateRandomKey(int size) {
        byte[] key = new byte[size];
        new java.security.SecureRandom().nextBytes(key);
        return key;
    }
    
    // Content encryption methods (same as previous example)
    private byte[] encryptWithContentKey(byte[] key, byte[] data) { /* ... */ }
    private byte[] decryptWithContentKey(byte[] key, byte[] encryptedData) { /* ... */ }
}

import com.azure.core.cryptography.*;
import java.time.Duration;
import java.util.concurrent.atomic.AtomicReference;

class RotatingKeyEncryptionService {
    private final AtomicReference<AsyncKeyEncryptionKey> currentKey = new AtomicReference<>();
    private final AsyncKeyEncryptionKeyResolver keyResolver;
    private final List<AsyncKeyEncryptionKey> keyRotationOrder;
    private volatile int currentKeyIndex = 0;
    
    public RotatingKeyEncryptionService(List<AsyncKeyEncryptionKey> keys, AsyncKeyEncryptionKeyResolver resolver) {
        this.keyRotationOrder = keys;
        this.keyResolver = resolver;
        this.currentKey.set(keys.get(0));
        
        // Start key rotation timer
        startKeyRotation();
    }
    
    public Mono<EncryptedBlob> encrypt(BinaryData data) {
        AsyncKeyEncryptionKey key = currentKey.get();
        return encryptWithKey(key, data);
    }
    
    public Mono<BinaryData> decrypt(EncryptedBlob encryptedBlob) {
        return keyResolver.resolveKey(encryptedBlob.getKeyId())
            .switchIfEmpty(Mono.error(new IllegalArgumentException("Key not found: " + encryptedBlob.getKeyId())))
            .flatMap(kek -> decryptWithKey(kek, encryptedBlob));
    }
    
    private void startKeyRotation() {
        // Rotate keys every hour (simplified example)
        Flux.interval(Duration.ofHours(1))
            .subscribe(tick -> rotateKey());
    }
    
    private void rotateKey() {
        currentKeyIndex = (currentKeyIndex + 1) % keyRotationOrder.size();
        currentKey.set(keyRotationOrder.get(currentKeyIndex));
        System.out.println("Rotated to key index: " + currentKeyIndex);
    }
    
    private Mono<EncryptedBlob> encryptWithKey(AsyncKeyEncryptionKey kek, BinaryData data) {
        // Implementation similar to previous examples
        return Mono.empty(); // Placeholder
    }
    
    private Mono<BinaryData> decryptWithKey(AsyncKeyEncryptionKey kek, EncryptedBlob blob) {
        // Implementation similar to previous examples  
        return Mono.empty(); // Placeholder
    }
}