Java 类com.amazonaws.encryptionsdk.AwsCrypto 实例源码

public KMSEncryptor(KMSManager kmsManager, AWSCredentialsProvider awsCredentials, ClientConfiguration clientConfiguration, SecretsGroupIdentifier groupIdentifier, AwsCrypto awsCrypto, EncryptionStrength encryptionStrength) {
    this.awsCredentials = awsCredentials;
    this.clientConfiguration = clientConfiguration;
    this.groupIdentifier = groupIdentifier;
    this.kmsManager = kmsManager;

    if (encryptionStrength.equals(EncryptionStrength.AES_128)) {
        awsCrypto.setEncryptionAlgorithm(CryptoAlgorithm.ALG_AES_128_GCM_IV12_TAG16_HKDF_SHA256_ECDSA_P256);
    } else if (encryptionStrength.equals(EncryptionStrength.AES_256)) {
        awsCrypto.setEncryptionAlgorithm(CryptoAlgorithm.ALG_AES_256_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384);
    } else {
        throw new IllegalArgumentException(String.format("Unrecognized encryption strength %s", encryptionStrength.toString()));
    }

    this.crypto = awsCrypto;
}

@BeforeMethod
public void setUp() throws Exception {
    AWSCredentialsProvider mockCredentials = mock(AWSCredentialsProvider.class);
    ClientConfiguration mockConfig = mock(ClientConfiguration.class);
    SecretsGroupIdentifier group = new SecretsGroupIdentifier(Region.US_WEST_1, "test.group");

    this.mockAwsCrypto = mock(AwsCrypto.class);
    this.mockKmsManager = mock(KMSManager.class);
    KMSEncryptor encryptor = new KMSEncryptor(mockKmsManager, mockCredentials, mockConfig, group, mockAwsCrypto, EncryptionStrength.AES_256);

    this.kmsEncryptor = spy(encryptor);
    this.mockProvider = mock(KmsMasterKeyProvider.class);
    doReturn(mockProvider).when(kmsEncryptor).getProvider();

    // Verify the expected encryption algorithm was set.
    verify(mockAwsCrypto, times(1)).setEncryptionAlgorithm(
            CryptoAlgorithm.ALG_AES_256_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384);
}

private static void escrowDecrypt(final String fileName) throws Exception {
    // You can decrypt the stream using only the private key.
    // This method does not call KMS.

    // 1. Instantiate the SDK
    final AwsCrypto crypto = new AwsCrypto();

    // 2. Instantiate a JCE master key provider
    // This method call uses the escrowed private key, not null 
    final JceMasterKey escrowPriv = JceMasterKey.getInstance(publicEscrowKey, privateEscrowKey, "Escrow", "Escrow",
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding");

    // 3. Decrypt the file
    // To simplify the code, we omit the encryption context. Production code should always 
    // use an encryption context. For an example, see the other SDK samples.
    final FileInputStream in = new FileInputStream(fileName + ".encrypted");
    final FileOutputStream out = new FileOutputStream(fileName + ".deescrowed");
    final CryptoOutputStream<?> decryptingStream = crypto.createDecryptingStream(escrowPriv, out);
    IOUtils.copy(in, decryptingStream);
    in.close();
    decryptingStream.close();

}

@Test
public void testMultipleKmsKeys() {
    final MockKMSClient kms = new MockKMSClient();
    final String arn1 = kms.createKey().getKeyMetadata().getArn();
    final String arn2 = kms.createKey().getKeyMetadata().getArn();
    MasterKeyProvider<KmsMasterKey> prov = legacyConstruct(kms, arn1, arn2);
    KmsMasterKey mk1 = prov.getMasterKey(arn1);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], KmsMasterKey> ct = crypto.encryptData(prov, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], KmsMasterKey> result = crypto.decryptData(prov, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));
}

@Test
public void testMixedKeys() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    final MockKMSClient kms = new MockKMSClient();
    final String arn2 = kms.createKey().getKeyMetadata().getArn();
    MasterKeyProvider<KmsMasterKey> prov = legacyConstruct(kms);
    KmsMasterKey mk2 = prov.getMasterKey(arn2);
    final MasterKeyProvider<?> mkp = MultipleProviderFactory.buildMultiProvider(mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], ?> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], ?> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    assertMultiReturnsKeys(mkp, mk1, mk2);
}

@Test
public void testMixedKeysSingleDecrypt() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    final MockKMSClient kms = new MockKMSClient();
    final String arn2 = kms.createKey().getKeyMetadata().getArn();
    MasterKeyProvider<KmsMasterKey> prov = legacyConstruct(kms);
    KmsMasterKey mk2 = prov.getMasterKey(arn2);
    final MasterKeyProvider<?> mkp = MultipleProviderFactory.buildMultiProvider(mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], ?> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    CryptoResult<byte[], ?> result = crypto.decryptData(mk1, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    result = crypto.decryptData(mk2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}

@Test
public void whenConstructedWithoutArguments_canUseMultipleRegions() throws Exception {
    KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder().build();

    for (String key : KMSTestFixtures.TEST_KEY_IDS) {
        byte[] ciphertext =
                new AwsCrypto().encryptData(
                        KmsMasterKeyProvider.builder()
                            .withKeysForEncryption(key)
                            .build(),
                        new byte[1]
                ).getResult();

        new AwsCrypto().decryptData(mkp, ciphertext);
    }
}

@SuppressWarnings("deprecation") @Test(expected = CannotUnwrapDataKeyException.class)
public void whenLegacyConstructorsUsed_multiRegionDecryptIsNotSupported() throws Exception {
    KmsMasterKeyProvider mkp = new KmsMasterKeyProvider();

    for (String key : KMSTestFixtures.TEST_KEY_IDS) {
        byte[] ciphertext =
                new AwsCrypto().encryptData(
                        KmsMasterKeyProvider.builder()
                                            .withKeysForEncryption(key)
                                            .build(),
                        new byte[1]
                ).getResult();

        new AwsCrypto().decryptData(mkp, ciphertext);
    }
}

@Test
public void whenHandlerConfigured_handlerIsInvoked() throws Exception {
    RequestHandler2 handler = spy(new RequestHandler2() {});
    KmsMasterKeyProvider mkp =
            KmsMasterKeyProvider.builder()
                                .withClientBuilder(
                                        AWSKMSClientBuilder.standard()
                                            .withRequestHandlers(handler)
                                )
                                .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
                                .build();

    new AwsCrypto().encryptData(mkp, new byte[1]);

    verify(handler).beforeRequest(any());
}

@Test
public void whenCustomCredentialsSet_theyAreUsed() throws Exception {
    AWSCredentialsProvider customProvider = spy(new DefaultAWSCredentialsProviderChain());

    KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder()
                                                   .withCredentials(customProvider)
                                                   .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
                                                   .build();

    new AwsCrypto().encryptData(mkp, new byte[1]);

    verify(customProvider, atLeastOnce()).getCredentials();

    AWSCredentials customCredentials = spy(customProvider.getCredentials());

    mkp = KmsMasterKeyProvider.builder()
                                                   .withCredentials(customCredentials)
                                                   .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
                                                   .build();

    new AwsCrypto().encryptData(mkp, new byte[1]);

    verify(customCredentials, atLeastOnce()).getAWSSecretKey();
}

@Test
public void testLegacyGrantTokenPassthrough() throws Exception {
    MockKMSClient client = spy(new MockKMSClient());

    String key1 = client.createKey().getKeyMetadata().getArn();

    KmsMasterKeyProvider mkp = new KmsMasterKeyProvider(client, getRegion(fromName("us-west-2")), singletonList(key1));

    mkp.addGrantToken("x");
    mkp.setGrantTokens(new ArrayList<>(Arrays.asList("y")));
    mkp.setGrantTokens(new ArrayList<>(Arrays.asList("a", "b")));
    mkp.addGrantToken("c");

    byte[] ciphertext = new AwsCrypto().encryptData(mkp, new byte[0]).getResult();

    ArgumentCaptor<GenerateDataKeyRequest> gdkr = ArgumentCaptor.forClass(GenerateDataKeyRequest.class);
    verify(client, times(1)).generateDataKey(gdkr.capture());

    List<String> grantTokens = gdkr.getValue().getGrantTokens();
    assertTrue(grantTokens.contains("a"));
    assertTrue(grantTokens.contains("b"));
    assertTrue(grantTokens.contains("c"));
    assertFalse(grantTokens.contains("x"));
    assertFalse(grantTokens.contains("z"));
}

@Test
public void multipleKeys() throws GeneralSecurityException {
    addEntry("key1");
    addEntry("key2");
    final KeyStoreProvider mkp = new KeyStoreProvider(ks, PP, "KeyStore", "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            "key1",
            "key2");
    @SuppressWarnings("unused")
    final JceMasterKey mk1 = mkp.getMasterKey("key1");
    final JceMasterKey mk2 = mkp.getMasterKey("key2");
    final AwsCrypto crypto = new AwsCrypto();
    final CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Order is non-deterministic
    assertEquals(1, result.getMasterKeys().size());

    // Delete the first key and see if it works
    ks.deleteEntry("key1");
    result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}

@Test
public void escrowAndSymmetricSecondProvider() throws GeneralSecurityException {
    addPublicEntry("key1");
    addEntry("key2");
    final KeyStoreProvider mkp = new KeyStoreProvider(ks, PP, "KeyStore", "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            "key1",
            "key2");
    @SuppressWarnings("unused")
    final JceMasterKey mk1 = mkp.getMasterKey("key1");
    final JceMasterKey mk2 = mkp.getMasterKey("key2");
    final AwsCrypto crypto = new AwsCrypto();
    final CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    final KeyStoreProvider mkp2 = new KeyStoreProvider(ks, PP, "KeyStore", "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            "key1");
    CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(mkp2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only could have decrypted with the keypair
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}

@Test
public void keystoreAndRawProvider() throws GeneralSecurityException, IOException {
    addEntry("key1");
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey jcep = JceMasterKey.getInstance(k1, "jce", "1", "AES/GCM/NoPadding");
    final KeyStoreProvider ksp = new KeyStoreProvider(ks, PP, "KeyStore", "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            "key1");

    MasterKeyProvider<JceMasterKey> multiProvider = MultipleProviderFactory.buildMultiProvider(JceMasterKey.class,
            jcep, ksp);

    assertEquals(jcep, multiProvider.getMasterKey("jce", "1"));

    final AwsCrypto crypto = new AwsCrypto();
    final CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(multiProvider, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(multiProvider, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(jcep, result.getMasterKeys().get(0));

    // Decrypt just using each individually
    assertArrayEquals(PLAINTEXT, crypto.decryptData(jcep, ct.getResult()).getResult());
    assertArrayEquals(PLAINTEXT, crypto.decryptData(ksp, ct.getResult()).getResult());
}

private byte[] getTestHeaders() {
    final CryptoAlgorithm cryptoAlgorithm_ = AwsCrypto.getDefaultCryptoAlgorithm();
    final int frameSize_ = AwsCrypto.getDefaultFrameSize();
    final Map<String, String> encryptionContext = Collections.<String, String> emptyMap();

    final EncryptionMaterialsRequest encryptionMaterialsRequest = EncryptionMaterialsRequest.newBuilder()
                                                                                            .setContext(encryptionContext)
                                                                                            .setRequestedAlgorithm(cryptoAlgorithm_)
                                                                                            .build();

    final EncryptionMaterials encryptionMaterials = new DefaultCryptoMaterialsManager(masterKeyProvider_)
            .getMaterialsForEncrypt(encryptionMaterialsRequest);

    final EncryptionHandler encryptionHandler = new EncryptionHandler(frameSize_, encryptionMaterials);

    // create the ciphertext headers by calling encryption handler.
    final byte[] in = new byte[0];
    final int ciphertextLen = encryptionHandler.estimateOutputSize(in.length);
    final byte[] ciphertext = new byte[ciphertextLen];
    encryptionHandler.processBytes(in, 0, in.length, ciphertext, 0);
    return ciphertext;
}

@Test(expected = BadCiphertextException.class)
public void tamperCiphertext() {
    final CryptoAlgorithm cryptoAlgorithm = AwsCrypto.getDefaultCryptoAlgorithm();
    final byte[] content = RandomBytesGenerator.generate(contentLen_);
    final byte[] keyBytes = RandomBytesGenerator.generate(cryptoAlgorithm.getKeyLength());
    final byte[] nonce = RandomBytesGenerator.generate(cryptoAlgorithm.getNonceLen());

    final SecretKey key = new SecretKeySpec(keyBytes, cryptoAlgorithm.name());
    CipherHandler cipherHandler = createCipherHandler(key, cryptoAlgorithm, Cipher.ENCRYPT_MODE);
    final byte[] encryptedBytes = cipherHandler.cipherData(nonce, contentAad_, content, 0, content.length);

    encryptedBytes[0] += 1; // tamper the first byte in ciphertext

    cipherHandler = createCipherHandler(key, cryptoAlgorithm, Cipher.DECRYPT_MODE);
    cipherHandler.cipherData(nonce, contentAad_, encryptedBytes, 0, encryptedBytes.length);
}

@Test
public void testMultipleJceKeys() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    final SecretKeySpec k2 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk2 = JceMasterKey.getInstance(k2, "jce", "2", WRAPPING_ALG);
    final MasterKeyProvider<JceMasterKey> mkp = MultipleProviderFactory.buildMultiProvider(JceMasterKey.class,
            mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    assertMultiReturnsKeys(mkp, mk1, mk2);
}

@Test
public void testMultipleJceKeysSingleDecrypt() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    final SecretKeySpec k2 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk2 = JceMasterKey.getInstance(k2, "jce", "2", WRAPPING_ALG);
    final MasterKeyProvider<JceMasterKey> mkp = MultipleProviderFactory.buildMultiProvider(JceMasterKey.class,
            mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(mk1, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    result = crypto.decryptData(mk2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}

@Test
public void testMixedKeys() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    StaticMasterKey mk2 = new StaticMasterKey("mock1");
    final MasterKeyProvider<?> mkp = MultipleProviderFactory.buildMultiProvider(mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], ?> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], ?> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    assertMultiReturnsKeys(mkp, mk1, mk2);
}

@Test
public void testMixedKeysSingleDecrypt() {
    final SecretKeySpec k1 = new SecretKeySpec(generate(32), "AES");
    final JceMasterKey mk1 = JceMasterKey.getInstance(k1, "jce", "1", WRAPPING_ALG);
    StaticMasterKey mk2 = new StaticMasterKey("mock1");

    final MasterKeyProvider<?> mkp = MultipleProviderFactory.buildMultiProvider(mk1, mk2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], ?> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    CryptoResult<byte[], ?> result = crypto.decryptData(mk1, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    result = crypto.decryptData(mk2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}

/**
 * {@inheritDoc }
 */
@Override
public void decryptFile(
        final String encryptedFilename, 
        final String decryptedFilename) {

    final KmsMasterKeyProvider provider
            = new KmsMasterKeyProvider(
                    new DefaultAWSCredentialsProviderChain());

    final AwsCrypto awsCrypto
            = new AwsCrypto();

    try (final FileInputStream fileInputStream
            = new FileInputStream(
                    encryptedFilename);

            final FileOutputStream fileOutputStream
                    = new FileOutputStream(
                            decryptedFilename);

            final CryptoInputStream<?> decryptingStream
                    = awsCrypto
                            .createDecryptingStream(
                                    provider, 
                                    fileInputStream)) {

        IOUtils.copy(
                decryptingStream,
                fileOutputStream);

    } catch (IOException exception) {
        throw new DecryptionException(exception);
    }
}

/**
 * {@inheritDoc }
 */
@Override
public String decryptFile(
        final String encryptedFilename) {

    final KmsMasterKeyProvider provider
            = new KmsMasterKeyProvider(
                    new DefaultAWSCredentialsProviderChain());

    final AwsCrypto awsCrypto
            = new AwsCrypto();

    try (final FileInputStream fileInputStream
            = new FileInputStream(
                    encryptedFilename);

            final CryptoInputStream<?> decryptingStream
                    = awsCrypto
                            .createDecryptingStream(
                                    provider, 
                                    fileInputStream)) {

        return IOUtils.toString(
                decryptingStream);

    } catch (IOException exception) {
        throw new DecryptionException(exception);
    }
}

public static KMSEncryptor fromCredentials(AWSCredentialsProvider awsCredentials,
                                           ClientConfiguration clientConfiguration,
                                           SecretsGroupIdentifier groupIdentifier,
                                           EncryptionStrength encryptionStrength) {
    KMSManager manager = KMSManager.fromCredentials(awsCredentials, clientConfiguration, groupIdentifier);
    return new KMSEncryptor(manager, awsCredentials, clientConfiguration, groupIdentifier, new AwsCrypto(), encryptionStrength);
}

@BeforeMethod
public void setUp() {
    mockCredentials = mock(AWSCredentialsProvider.class);
    mockClient = mock(AmazonIdentityManagementClient.class);
    ClientConfiguration mockConfig = mock(ClientConfiguration.class);
    IAMPolicyManager policyManager = new IAMPolicyManager(mockClient, mockCredentials, mockConfig);

    // The mockito spy acts like original object but mocks out the getAccount() method. As the getAccount() calls
    // directly rather than via a client that we can pass in we need to mock this out using a spy.
    partiallyMockedPolicyManager = spy(policyManager);
    doReturn(ACCOUNT).when(partiallyMockedPolicyManager).getAccount();

    // Set up KMSEncryptor for testing the policy creation methods. This gets a bit complicated but we need to
    // mock all the AWS dependencies from the KMSManager before using it to create the KMSEncryptor. The getAliasArn
    // needs to be mocked out with a spy to stop the call to getAccount.
    mockKMSClient = mock(AWSKMSClient.class);
    KMSManager kmsManager = new KMSManager(mockKMSClient, mockCredentials, mockConfig, group);
    KMSManager partiallyMockedKMSManager = spy(kmsManager);
    doReturn(KMS_ALIAS_ARN).when(partiallyMockedKMSManager).getAliasArn();
    kmsEncryptor = new KMSEncryptor(partiallyMockedKMSManager, mockCredentials, mockConfig, group, mock(AwsCrypto.class), EncryptionStrength.AES_256);

    // Set up store for testing the policy creation methods. Mock out the getArn method with a spy to stop the
    // call to getAccount().
    mockDynamoDBClient = mock(AmazonDynamoDBClient.class);
    DynamoDB store = new DynamoDB(mockDynamoDBClient, mockCredentials, mockConfig, group, new ReentrantReadWriteLock());
    partiallyMockedStore = spy(store);
    doReturn(DYNAMODB_ARN).when(partiallyMockedStore).getArn();
}

public static void main(String[] args) throws IOException {
    srcFile = args[0];

    // In this example, we generate a random key. In practice, 
    // you would get a key from an existing store
    SecretKey cryptoKey = retrieveEncryptionKey();

    // Create a JCE master key provider using the random key and an AES-GCM encryption algorithm
    JceMasterKey masterKey = JceMasterKey.getInstance(cryptoKey, "Example", "RandomKey", "AES/GCM/NoPadding");

    // Instantiate the SDK
    AwsCrypto crypto = new AwsCrypto();

    // Create an encryption context to identify this ciphertext
    Map<String, String> context = Collections.singletonMap("Example", "FileStreaming");

    // Because the file might be to large to load into memory, we stream the data, instead of 
    //loading it all at once.
    FileInputStream in = new FileInputStream(srcFile);
    CryptoInputStream<JceMasterKey> encryptingStream = crypto.createEncryptingStream(masterKey, in, context);

    FileOutputStream out = new FileOutputStream(srcFile + ".encrypted");
    IOUtils.copy(encryptingStream, out);
    encryptingStream.close();
    out.close();

    // Decrypt the file. Verify the encryption context before returning the plaintext.
    in = new FileInputStream(srcFile + ".encrypted");
    CryptoInputStream<JceMasterKey> decryptingStream = crypto.createDecryptingStream(masterKey, in);
    // Does it contain the expected encryption context?
    if (!"FileStreaming".equals(decryptingStream.getCryptoResult().getEncryptionContext().get("Example"))) {
        throw new IllegalStateException("Bad encryption context");
    }

    // Return the plaintext data
    out = new FileOutputStream(srcFile + ".decrypted");
    IOUtils.copy(decryptingStream, out);
    decryptingStream.close();
    out.close();
}

private static void standardEncrypt(final String kmsArn, final String fileName) throws Exception {
    // Encrypt with the KMS CMK and the escrowed public key
    // 1. Instantiate the SDK
    final AwsCrypto crypto = new AwsCrypto();

    // 2. Instantiate a KMS master key provider
    final KmsMasterKeyProvider kms = new KmsMasterKeyProvider(kmsArn);

    // 3. Instantiate a JCE master key provider
    // Because the user does not have access to the private escrow key,
    // they pass in "null" for the private key parameter.
    final JceMasterKey escrowPub = JceMasterKey.getInstance(publicEscrowKey, null, "Escrow", "Escrow",
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding");

    // 4. Combine the providers into a single master key provider
    final MasterKeyProvider<?> provider = MultipleProviderFactory.buildMultiProvider(kms, escrowPub);

    // 5. Encrypt the file
    // To simplify the code, we omit the encryption context. Production code should always 
    // use an encryption context. For an example, see the other SDK samples.
    final FileInputStream in = new FileInputStream(fileName);
    final FileOutputStream out = new FileOutputStream(fileName + ".encrypted");
    final CryptoOutputStream<?> encryptingStream = crypto.createEncryptingStream(provider, out);

    IOUtils.copy(in, encryptingStream);
    in.close();
    encryptingStream.close();
}

private static void standardDecrypt(final String kmsArn, final String fileName) throws Exception {
    // Decrypt with the KMS CMK and the escrow public key. You can use a combined provider, 
    // as shown here, or just the KMS master key provider.

    // 1. Instantiate the SDK
    final AwsCrypto crypto = new AwsCrypto();

    // 2. Instantiate a KMS master key provider
    final KmsMasterKeyProvider kms = new KmsMasterKeyProvider(kmsArn);

    // 3. Instantiate a JCE master key provider
    // Because the user does not have access to the private 
    // escrow key, they pass in "null" for the private key parameter.
    final JceMasterKey escrowPub = JceMasterKey.getInstance(publicEscrowKey, null, "Escrow", "Escrow",
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding");

    // 4. Combine the providers into a single master key provider
    final MasterKeyProvider<?> provider = MultipleProviderFactory.buildMultiProvider(kms, escrowPub);

    // 5. Decrypt the file
    // To simplify the code, we omit the encryption context. Production code should always 
    // use an encryption context. For an example, see the other SDK samples.
    final FileInputStream in = new FileInputStream(fileName + ".encrypted");
    final FileOutputStream out = new FileOutputStream(fileName + ".decrypted");
    final CryptoOutputStream<?> decryptingStream = crypto.createDecryptingStream(provider, out);
    IOUtils.copy(in, decryptingStream);
    in.close();
    decryptingStream.close();
}

@Test
public void testMultipleKmsKeysSingleDecrypt() {
    final MockKMSClient kms = new MockKMSClient();
    final String arn1 = kms.createKey().getKeyMetadata().getArn();
    final String arn2 = kms.createKey().getKeyMetadata().getArn();
    MasterKeyProvider<KmsMasterKey> prov = legacyConstruct(kms, arn1, arn2);
    KmsMasterKey mk1 = prov.getMasterKey(arn1);
    KmsMasterKey mk2 = prov.getMasterKey(arn2);

    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], KmsMasterKey> ct = crypto.encryptData(prov, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    CryptoResult<byte[], KmsMasterKey> result = crypto.decryptData(mk1, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    result = crypto.decryptData(mk2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));

    // Delete one of the two keys and ensure it's still decryptable
    kms.deleteKey(arn1);

    result = crypto.decryptData(prov, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}

@Test
public void testMultipleRegionKmsKeys() {
    final MockKMSClient us_east_1 = new MockKMSClient();
    us_east_1.setRegion(Region.getRegion(Regions.US_EAST_1));
    final MockKMSClient eu_west_1 = new MockKMSClient();
    eu_west_1.setRegion(Region.getRegion(Regions.EU_WEST_1));
    final String arn1 = us_east_1.createKey().getKeyMetadata().getArn();
    final String arn2 = eu_west_1.createKey().getKeyMetadata().getArn();
    KmsMasterKeyProvider provE = legacyConstruct(us_east_1, Region.getRegion(Regions.US_EAST_1));
    KmsMasterKeyProvider provW = legacyConstruct(eu_west_1, Region.getRegion(Regions.EU_WEST_1));
    KmsMasterKey mk1 = provE.getMasterKey(arn1);
    KmsMasterKey mk2 = provW.getMasterKey(arn2);

    final MasterKeyProvider<KmsMasterKey> mkp = MultipleProviderFactory.buildMultiProvider(KmsMasterKey.class,
                                                                                           mk1, mk2);
    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], KmsMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    CryptoResult<byte[], KmsMasterKey> result = crypto.decryptData(mk1, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    result = crypto.decryptData(mk2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));

    assertMultiReturnsKeys(mkp, mk1, mk2);

    // Delete one of the two keys and ensure it's still decryptable
    us_east_1.deleteKey(arn1);

    result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}

@Test
public void whenShortTimeoutSet_timesOut() throws Exception {
    // By setting a timeout of 1ms, it's not physically possible to complete both the us-west-2 and eu-central-1
    // requests due to speed of light limits.
    KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder()
                                                   .withClientBuilder(
                                                           AWSKMSClientBuilder.standard()
                                                            .withClientConfiguration(
                                                                    new ClientConfiguration()
                                                                        .withRequestTimeout(1)
                                                            )
                                                   )
                                                   .withKeysForEncryption(Arrays.asList(KMSTestFixtures.TEST_KEY_IDS))
                                                   .build();

    try {
        new AwsCrypto().encryptData(mkp, new byte[1]);
        fail("Expected exception");
    } catch (Exception e) {
        if (e instanceof AbortedException) {
            // ok - one manifestation of a timeout
        } else if (e.getCause() instanceof HttpRequestTimeoutException) {
            // ok - another kind of timeout
        } else {
            throw e;
        }
    }
}

@Test
public void whenBuilderCloned_credentialsAndConfigurationAreRetained() throws Exception {
    AWSCredentialsProvider customProvider1 = spy(new DefaultAWSCredentialsProviderChain());
    AWSCredentialsProvider customProvider2 = spy(new DefaultAWSCredentialsProviderChain());

    KmsMasterKeyProvider.Builder builder = KmsMasterKeyProvider.builder()
            .withCredentials(customProvider1)
            .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0]);

    KmsMasterKeyProvider.Builder builder2 = builder.clone();

    // This will mutate the first builder to add the new key and change the creds, but leave the clone unchanged.
    MasterKeyProvider<?> mkp2 = builder.withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[1]).withCredentials(customProvider2).build();
    MasterKeyProvider<?> mkp1 = builder2.build();

    CryptoResult<byte[], ?> result = new AwsCrypto().encryptData(mkp1, new byte[0]);

    assertEquals(KMSTestFixtures.TEST_KEY_IDS[0], result.getMasterKeyIds().get(0));
    assertEquals(1, result.getMasterKeyIds().size());
    verify(customProvider1, atLeastOnce()).getCredentials();
    verify(customProvider2, never()).getCredentials();

    reset(customProvider1, customProvider2);

    result = new AwsCrypto().encryptData(mkp2, new byte[0]);

    assertTrue(result.getMasterKeyIds().contains(KMSTestFixtures.TEST_KEY_IDS[0]));
    assertTrue(result.getMasterKeyIds().contains(KMSTestFixtures.TEST_KEY_IDS[1]));
    assertEquals(2, result.getMasterKeyIds().size());
    verify(customProvider1, never()).getCredentials();
    verify(customProvider2, atLeastOnce()).getCredentials();
}

@Test
public void whenBuilderCloned_clientBuilderCustomizationIsRetained() throws Exception {
    RequestHandler2 handler = spy(new RequestHandler2() {});

    KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder()
            .withClientBuilder(
                    AWSKMSClientBuilder.standard().withRequestHandlers(handler)
            )
            .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
            .clone().build();

    new AwsCrypto().encryptData(mkp, new byte[0]);

    verify(handler, atLeastOnce()).beforeRequest(any());
}

@Test
public void testDecryptFromFile() throws Exception {
    AwsCrypto crypto = new AwsCrypto();
    final KmsMasterKeyProvider masterKeyProvider = new KmsMasterKeyProvider(kmsKeyId);
    byte ciphertextBytes[] = Files.readAllBytes(Paths.get(ciphertextFileName));
    byte plaintextBytes[] = Files.readAllBytes(Paths.get(plaintextFileName));
    final CryptoResult decryptResult = crypto.decryptData(
        masterKeyProvider,
        ciphertextBytes
    );
    assertArrayEquals(plaintextBytes, (byte[])decryptResult.getResult());
}

@Test
public void singleKeyPkcs1() throws GeneralSecurityException {
    addEntry("key1");
    final KeyStoreProvider mkp = new KeyStoreProvider(ks, PP, "KeyStore", "RSA/ECB/PKCS1Padding", "key1");
    final JceMasterKey mk1 = mkp.getMasterKey("key1");
    final AwsCrypto crypto = new AwsCrypto();
    final CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(1, ct.getMasterKeyIds().size());
    final CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));
}

@Test
public void singleKeyOaepSha1() throws GeneralSecurityException {
    addEntry("key1");
    final KeyStoreProvider mkp = new KeyStoreProvider(ks, PP, "KeyStore", "RSA/ECB/OAEPWithSHA-1AndMGF1Padding",
            "key1");
    final JceMasterKey mk1 = mkp.getMasterKey("key1");
    final AwsCrypto crypto = new AwsCrypto();
    final CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(1, ct.getMasterKeyIds().size());
    final CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));
}

@Test
public void singleKeyOaepSha256() throws GeneralSecurityException {
    addEntry("key1");
    final KeyStoreProvider mkp = new KeyStoreProvider(ks, PP, "KeyStore", "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            "key1");
    final JceMasterKey mk1 = mkp.getMasterKey("key1");
    final AwsCrypto crypto = new AwsCrypto();
    final CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(1, ct.getMasterKeyIds().size());
    final CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));
}

@Test(expected = CannotUnwrapDataKeyException.class)
public void encryptOnly() throws GeneralSecurityException {
    addPublicEntry("key1");
    final KeyStoreProvider mkp = new KeyStoreProvider(ks, PP, "KeyStore", "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            "key1");
    final AwsCrypto crypto = new AwsCrypto();
    final CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(1, ct.getMasterKeyIds().size());
    crypto.decryptData(mkp, ct.getResult());
}

@Test
public void escrowAndSymmetric() throws GeneralSecurityException {
    addPublicEntry("key1");
    addEntry("key2");
    final KeyStoreProvider mkp = new KeyStoreProvider(ks, PP, "KeyStore", "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            "key1",
            "key2");
    @SuppressWarnings("unused")
    final JceMasterKey mk1 = mkp.getMasterKey("key1");
    final JceMasterKey mk2 = mkp.getMasterKey("key2");
    final AwsCrypto crypto = new AwsCrypto();
    final CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());
    CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only could have decrypted with the keypair
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));

    // Delete the first key and see if it works
    ks.deleteEntry("key1");
    result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}

@Test
public void escrowCase() throws GeneralSecurityException, IOException {
    addEntry("escrowKey");
    KeyStore ks2 = KeyStore.getInstance(KeyStore.getDefaultType());
    ks2.load(null, PASSWORD);
    copyPublicPart(ks, ks2, "escrowKey");

    final KeyStoreProvider mkp = new KeyStoreProvider(ks, PP, "KeyStore", "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            "escrowKey");
    final KeyStoreProvider escrowProvider = new KeyStoreProvider(ks2, PP, "KeyStore",
            "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            "escrowKey");

    final JceMasterKey mk1 = escrowProvider.getMasterKey("escrowKey");
    final AwsCrypto crypto = new AwsCrypto();
    final CryptoResult<byte[], JceMasterKey> ct = crypto.encryptData(escrowProvider, PLAINTEXT);
    assertEquals(1, ct.getMasterKeyIds().size());

    try {
        crypto.decryptData(escrowProvider, ct.getResult());
        fail("Expected CannotUnwrapDataKeyException");
    } catch (final CannotUnwrapDataKeyException ex) {
        // expected
    }
    CryptoResult<byte[], JceMasterKey> result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only could have decrypted with the keypair
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));
}

/**
 * {@inheritDoc }
 */
@Override
public File encryptFile(final File file) {
    if (!file.exists())
        throw new IllegalArgumentException(
                String.format(
                        FILE_DOES_NOT_EXIST_PATTERN,
                        file.getAbsolutePath()));

    if (!file.isFile())
        throw new IllegalArgumentException(
                String.format(
                        FILE_IS_NOT_A_NORMAL_FILE_PATTERN,
                        file.getAbsolutePath()));

    if (!file.canRead())
        throw new IllegalArgumentException(
                String.format(
                    FILE_IS_NOT_READABLE_PATTERN, 
                    file.getAbsolutePath()));

    if (!file.getParentFile().canWrite()) 
        throw new IllegalArgumentException(
                String.format(
                        DIRECTORY_IS_NOT_READABLE,
                        file.getParentFile().getAbsolutePath()));

    File encryptedFile
            = new File(
                    new StringBuilder()
                            .append(file.getAbsoluteFile())
                            .append(SUFFIX)
                            .toString());

    final AwsCrypto awsCrypto
            = new AwsCrypto();

    try (final FileInputStream fileInputStream
            = new FileInputStream(file);

            final FileOutputStream fileOutputStram 
                    = new FileOutputStream(
                            encryptedFile);

            final CryptoOutputStream<?> encryptingStream
                    = awsCrypto
                           .createEncryptingStream(
                                   masterKeyProvider(),
                                   fileOutputStram)) {

        IOUtils.copy(
                fileInputStream,
                encryptingStream);

    } catch (IOException exception) {
        throw new EncryptionException(exception);
    } 

    return encryptedFile;
}