com.amazonaws.encryptionsdk.CryptoResult Java Examples

@Override
public byte[] decrypt(byte[] ciphertext, EncryptionContext context) {
    try {
        final CryptoResult<byte[], KmsMasterKey> decryptResult = crypto.decryptData(getProvider(), ciphertext);

        verify(decryptResult, context);

        return decryptResult.getResult();
    } catch (AwsCryptoException e) {
        if (isInvalidKeyException(e)) {
            throw new UnlimitedEncryptionNotSetException();
        } else {
            throw e;
        }
    }
}
 

@Test
public void escrowAndSymmetricSecondProvider() throws GeneralSecurityException, IOException {
    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());
}
 

@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 multipleKeys() throws Exception {
    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 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));
}
 

@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 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);
}
 

/**
 * Returns a decrypted {@link Field} with its original type preserved when type information
 * has been preserved in the AAD.
 *
 * @param result of the decryption operation
 * @return {@link Field}
 */
public Field createResultFieldDecrypt(CryptoResult<byte[], ?> result) {
  Field.Type fieldType = Field.Type.valueOf(
      result.getEncryptionContext().getOrDefault(SDC_FIELD_TYPE, Field.Type.BYTE_ARRAY.name())
  );

  // Field API prohibits STRING to BYTE_ARRAY conversion so this is a special case
  if (fieldType == Field.Type.BYTE_ARRAY) {
    return Field.create(result.getResult());
  }

  // Field API supports STRING to other primitive types.
  return Field.create(
      fieldType,
      new String(result.getResult())
  );
}
 

@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));
}
 

@Override
public void process(FieldBatch batch) throws StageException {
  Map<String, String> encryptionContext = new HashMap<>(conf.context);

  while(batch.next()) {
    Field field = batch.getField();
    // reviewer requested no use of Java 8 streams
    if (field != null && field.getValue() != null) { // process if field is present and non-null

      Optional<Field> input = encrypter.checkInputEncrypt(field);

      if (input.isPresent()) {
        byte[] bytes = prepare.apply(input.get(), encryptionContext);
        CryptoResult<byte[], ?> result = encrypter.process(bytes, encryptionContext);
        Field encryptedField = createResultField.apply(result);
        batch.replace(encryptedField);
      } else {
        return; // record sent to error, done with this record.
      }
    }
  }
}
 

/**
 * Test encryption and decryption.
 */
@Test
public void testEncryptDecrypt() {
    String encKey = "12345";
    byte[] testData = "test string".getBytes(StandardCharsets.UTF_8);
    byte[] encTestData = "enc test string".getBytes(StandardCharsets.UTF_8);

    AwsKmsEncryptionService awsKmsEncryptionSvc = Mockito.spy(new AwsKmsEncryptionService());
    awsKmsEncryptionSvc.setKeyId(encKey)
        .setCredentials(new BasicAWSCredentials("dummy", "dummy"))
        .setRegion(Region.getRegion(Regions.AP_SOUTH_1));

    AwsCrypto awsCrypto = Mockito.mock(AwsCrypto.class);
    KmsMasterKeyProvider prov = Mockito.mock(KmsMasterKeyProvider.class);
    CryptoResult encCryptoRes = Mockito.mock(CryptoResult.class);
    CryptoResult decCryptoRes = Mockito.mock(CryptoResult.class);

    Mockito.doReturn(awsCrypto).when(awsKmsEncryptionSvc).createClient();
    Mockito.doReturn(prov).when(awsKmsEncryptionSvc).createKmsMasterKeyProvider();

    awsKmsEncryptionSvc.init();

    Mockito.doReturn(encCryptoRes).when(awsCrypto).encryptData(prov, testData);
    Mockito.doReturn(encTestData).when(encCryptoRes).getResult();

    Mockito.doReturn(decCryptoRes).when(awsCrypto).decryptData(prov, encTestData);
    Mockito.doReturn(Arrays.asList(encKey)).when(decCryptoRes).getMasterKeyIds();
    Mockito.doReturn(testData).when(decCryptoRes).getResult();

    byte[] encData = awsKmsEncryptionSvc.encrypt(testData);
    byte[] actualOutput = awsKmsEncryptionSvc.decrypt(encData);

    Assert.assertArrayEquals(testData, actualOutput);
}
 

/** {@inheritDoc} */
@Override public byte[] decrypt(byte[] data) {
    if (crypto == null || prov == null)
        throw new IgniteException("The init() method was not called.");

    CryptoResult<byte[], KmsMasterKey> decryptRes = crypto.decryptData(prov, data);

    List<String> keyIds = decryptRes.getMasterKeyIds();

    if (keyIds != null && !keyIds.contains(keyId))
        throw new IgniteException("Wrong KMS key ID!");

    return decryptRes.getResult();
}
 

@Override
public CryptoResult<byte[], ?> process(byte[] in, Map<String, String> context) {
  CryptoResult<byte[], ?> result;
  if (mode == EncryptionMode.ENCRYPT) {
    result = crypto.encryptData(cmManager, in, context);
  } else {
    result = crypto.decryptData(cmManager, in);
  }
  return result;
}
 

@Override
public String decrypt(String ciphertext, EncryptionContext context) {
    final CryptoResult<String, KmsMasterKey> decryptResult = crypto.decryptString(getProvider(), ciphertext);

    verify(decryptResult, context);

    return decryptResult.getResult();
}
 

@Override
protected void process(Record record, SingleLaneBatchMaker singleLaneBatchMaker) throws StageException {
  Map<String, String> encryptionContext = new HashMap<>(conf.context);

  try {
    for (String fieldPath : conf.fieldPaths) {

      Field field = record.get(fieldPath);

      // reviewer requested no use of Java 8 streams
      if (field != null && field.getValue() != null) { // process if field is present and non-null

        Optional<Field> input = checkInput.apply(record, field);

        if (input.isPresent()) {
          byte[] bytes = prepare.apply(input.get(), encryptionContext);
          CryptoResult<byte[], ?> result = encrypter.process(bytes, encryptionContext);
          field = createResultField.apply(result);
          record.set(fieldPath, field);
        } else {
          return; // record sent to error, done with this record.
        }
      }
    }
  // The encryption process can throw a lot of unchecked exceptions that if not caught would terminate the pipeline
  } catch (Exception e) {
    Throwables.propagateIfPossible(e, StageException.class);
    throw new StageException(CryptoErrors.CRYPTO_07, e.toString(), e);
  }

  singleLaneBatchMaker.addRecord(record);
}
 

@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));
}
 

@Test
public void escrowAndSymmetric() throws Exception {
    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(expected = CannotUnwrapDataKeyException.class)
public void encryptOnly() throws Exception {
    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 singleKeyOaepSha256() throws Exception {
    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
public void singleKeyOaepSha1() throws Exception {
    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 singleKeyPkcs1() throws Exception {
    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 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 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 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 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 testEncrypt() throws Exception {
    String plaintext = "jsonblob";
    EncryptionContext mockContext = mock(EncryptionContext.class);
    CryptoResult mockCryptoResult = mock(CryptoResult.class);
    Map<String, String> contextMap = new HashMap<>();

    when(mockContext.toMap()).thenReturn(contextMap);
    when(mockCryptoResult.getResult()).thenReturn(encryptedPayload);
    when(mockAwsCrypto.encryptString(mockProvider, plaintext, contextMap)).thenReturn(
            mockCryptoResult);
    assertEquals(kmsEncryptor.encrypt(plaintext, mockContext), encryptedPayload);
}
 

private void verify(CryptoResult<?, KmsMasterKey> decryptResult, EncryptionContext context) {
    if (!decryptResult.getMasterKeyIds().get(0).equals(getKeyArn())) {
        throw new IllegalStateException("Wrong key id!");
    }

    for (final Map.Entry<String, String> e : context.toMap().entrySet()) {
        if (!e.getValue().equals(decryptResult.getEncryptionContext().get(e.getKey()))) {
            throw new IllegalStateException("Wrong Encryption Context!");
        }
    }
}