com.amazonaws.services.kms.model.EncryptRequest Java Examples

@Test
void testAsymmetricEncryptionIsNotAvailable(CapturedOutput output) {
    doThrow(InvalidKeyUsageException.class).when(mockKms).encrypt(any(EncryptRequest.class));

    try {
        // Asymmetric algorithm is not available, because an outdated AWS SDK is used. The textEncryptor will
        // print a warning and fall back to symmetric algorithm.
        // Trying to use an asymmetric key with the symmetric algorithm will lead to an exception.
        textEncryptor.encrypt(PLAINTEXT);
        failBecauseExceptionWasNotThrown(InvalidKeyUsageException.class);
    } catch (InvalidKeyUsageException ignored) {
        assertThat(output).contains(VERSION_HINT);
        final EncryptRequest expectedRequest = new EncryptRequest()
                .withKeyId("an-asymmetric-key")
                .withPlaintext(ByteBuffer.wrap(PLAINTEXT.getBytes()));
        verify(mockKms).encrypt(eq(expectedRequest));
    }
}
 

@Override
public GenerateDataKeyResult generateDataKey(GenerateDataKeyRequest req)
        throws AmazonServiceException, AmazonClientException {
    byte[] pt;
    if (req.getKeySpec() != null) {
        if (req.getKeySpec().contains("256")) {
            pt = new byte[32];
        } else if (req.getKeySpec().contains("128")) {
            pt = new byte[16];
        } else {
            throw new UnsupportedOperationException();
        }
    } else {
        pt = new byte[req.getNumberOfBytes()];
    }
    rnd.nextBytes(pt);
    ByteBuffer ptBuff = ByteBuffer.wrap(pt);
    EncryptResult encryptResult = encrypt(new EncryptRequest().withKeyId(req.getKeyId())
            .withPlaintext(ptBuff).withEncryptionContext(req.getEncryptionContext()));
    return new GenerateDataKeyResult().withKeyId(req.getKeyId())
            .withCiphertextBlob(encryptResult.getCiphertextBlob()).withPlaintext(ptBuff);

}
 

@Before
public void setUp() {
    mockKms = mock(AWSKMS.class);
    textEncryptor = new KmsTextEncryptor(mockKms, KMS_KEY_ID, SYMMETRIC_DEFAULT.toString());

    expectedEncryptRequest = new EncryptRequest();
    expectedEncryptRequest.setKeyId(KMS_KEY_ID);
    expectedEncryptRequest.setPlaintext(wrap(PLAINTEXT.getBytes()));
    expectedEncryptRequest.setEncryptionAlgorithm(SYMMETRIC_DEFAULT.toString());

    encryptResult = new EncryptResult();
    encryptResult.setCiphertextBlob(wrap(CIPHER_TEXT.getBytes()));
    when(mockKms.encrypt(any(EncryptRequest.class))).thenReturn(encryptResult);

    expectedDecryptRequest = new DecryptRequest();
    expectedDecryptRequest.setCiphertextBlob(wrap(CIPHER_TEXT.getBytes()));
    expectedDecryptRequest.setEncryptionAlgorithm(SYMMETRIC_DEFAULT.toString());

    decryptResult = new DecryptResult();
    decryptResult.setPlaintext(wrap(PLAINTEXT.getBytes()));
    when(mockKms.decrypt(any(DecryptRequest.class))).thenReturn(decryptResult);
}
 

@Override
public String encrypt(final String text) {
    Assert.hasText(kmsKeyId, "kmsKeyId must not be blank");
    if (text == null || text.isEmpty()) {
        return EMPTY_STRING;
    } else {
        final EncryptRequest encryptRequest = new EncryptRequest()
                .withKeyId(kmsKeyId)
                .withPlaintext(ByteBuffer.wrap(text.getBytes()));

        checkAlgorithm(encryptionAlgorithm);

        if (IS_ALGORITHM_AVAILABLE) {
            encryptRequest.setEncryptionAlgorithm(encryptionAlgorithm);
        }

        final ByteBuffer encryptedBytes = kms.encrypt(encryptRequest).getCiphertextBlob();

        return extractString(encryptedBytes, BASE64);
    }
}
 

@Test
void testDecryptEndpoint() {
    final String cipherText = Base64.getEncoder().encodeToString("cIpHeR".getBytes());

    // Config Server does a "test" encrypt with the given key
    doAnswer(invocation -> new EncryptResult().withCiphertextBlob(ByteBuffer.wrap(cipherText.getBytes())))
            .when(mockKms).encrypt(any(EncryptRequest.class));

    final ResponseEntity<String> response = rest.exchange(
            post(URI.create("/decrypt"))
                    .contentType(APPLICATION_FORM_URLENCODED)
                    .body(cipherText),
            String.class);
    assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
    assertThat(response.getBody()).isEqualTo("Hello World");
}
 

@Test
void testEncryptEndpoint() {
    final String plainText = "some-plaintext";
    final String cipherText = "cIpHeR";

    doAnswer(invocation -> new EncryptResult().withCiphertextBlob(ByteBuffer.wrap(cipherText.getBytes())))
            .when(mockKms).encrypt(any(EncryptRequest.class));

    final ResponseEntity<String> response = rest.exchange(
            post(URI.create("/encrypt"))
                    .contentType(APPLICATION_FORM_URLENCODED)
                    .body(plainText),
            String.class);
    assertThat(response.getStatusCode()).isEqualTo(HttpStatus.OK);
    assertThat(response.getBody()).isEqualTo(Base64.getEncoder().encodeToString(cipherText.getBytes()));
}
 

@Test
void testEncrypt() {
    final byte[] cipherTextBytes = "bla".getBytes();
    final String expectedCipherString = Base64.getEncoder().encodeToString(cipherTextBytes);
    doReturn(new EncryptResult().withCiphertextBlob(ByteBuffer.wrap(cipherTextBytes)))
            .when(mockKms).encrypt(any(EncryptRequest.class));

    final String mySecret = "my-secret";
    final String encryptedString = textEncryptor.encrypt(mySecret);
    assertThat(encryptedString).isEqualTo(expectedCipherString);

    final EncryptRequest encryptRequest = new EncryptRequest()
            .withEncryptionAlgorithm("RSAES_OAEP_SHA_1")
            .withKeyId("asymmetric-sha1-sample-key")
            .withPlaintext(ByteBuffer.wrap(mySecret.getBytes()));
    verify(mockKms).encrypt(eq(encryptRequest));
}
 

@Override
public Observable<Encrypted> encrypt(VertxContext<Server> vertxContext, byte[] plainBytes) {
    SfsVertx sfsVertx = vertxContext.vertx();
    return Observable.defer(() -> {
        byte[] cloned = Arrays.copyOf(plainBytes, plainBytes.length);
        return RxHelper.executeBlocking(sfsVertx.getOrCreateContext(), sfsVertx.getBackgroundPool(), () -> {
            try {
                EncryptRequest req =
                        new EncryptRequest()
                                .withKeyId(keyId)
                                .withPlaintext(ByteBuffer.wrap(cloned));
                ByteBuffer buffer = kms.encrypt(req).getCiphertextBlob();
                byte[] b = new byte[buffer.remaining()];
                buffer.get(b);
                return new Encrypted(b, String.format("xppsaws:%s", keyId));
            } finally {
                Arrays.fill(cloned, (byte) 0);
            }
        });
    });
}
 

@Override
public GenerateDataKeyResult generateDataKey(GenerateDataKeyRequest req) throws AmazonServiceException,
        AmazonClientException {
    byte[] pt;
    if (req.getKeySpec() != null) {
        if (req.getKeySpec().contains("256")) {
            pt = new byte[32];
        } else if (req.getKeySpec().contains("128")) {
            pt = new byte[16];
        } else {
            throw new java.lang.UnsupportedOperationException();
        }
    } else {
        pt = new byte[req.getNumberOfBytes()];
    }
    rnd.nextBytes(pt);
    ByteBuffer ptBuff = ByteBuffer.wrap(pt);
    EncryptResult encryptResult = encrypt0(new EncryptRequest().withKeyId(req.getKeyId()).withPlaintext(ptBuff)
            .withEncryptionContext(req.getEncryptionContext()));
    String arn = retrieveArn(req.getKeyId());
    return new GenerateDataKeyResult().withKeyId(arn).withCiphertextBlob(encryptResult.getCiphertextBlob())
            .withPlaintext(ptBuff);
}
 

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

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

    KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder()
                                                   .withKeysForEncryption(key1, key2)
                                                   .withCustomClientFactory(ignored -> client)
                                                   .build();

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

    ArgumentCaptor<GenerateDataKeyRequest> gdkr = ArgumentCaptor.forClass(GenerateDataKeyRequest.class);
    verify(client, times(1)).generateDataKey(gdkr.capture());
    assertTrue(getUA(gdkr.getValue()).contains(VersionInfo.USER_AGENT));

    ArgumentCaptor<EncryptRequest> encr = ArgumentCaptor.forClass(EncryptRequest.class);
    verify(client, times(1)).encrypt(encr.capture());
    assertTrue(getUA(encr.getValue()).contains(VersionInfo.USER_AGENT));

    ArgumentCaptor<DecryptRequest> decr = ArgumentCaptor.forClass(DecryptRequest.class);
    verify(client, times(1)).decrypt(decr.capture());
    assertTrue(getUA(decr.getValue()).contains(VersionInfo.USER_AGENT));
}
 

@Override
public DataKey<KmsMasterKey> encryptDataKey(final CryptoAlgorithm algorithm,
        final Map<String, String> encryptionContext,
        final DataKey<?> dataKey) {
    final SecretKey key = dataKey.getKey();
    if (!key.getFormat().equals("RAW")) {
        throw new IllegalArgumentException("Only RAW encoded keys are supported");
    }
    try {
        final EncryptResult encryptResult = kms_.get().encrypt(updateUserAgent(
                new EncryptRequest()
                        .withKeyId(id_)
                        .withPlaintext(ByteBuffer.wrap(key.getEncoded()))
                        .withEncryptionContext(encryptionContext)
                        .withGrantTokens(grantTokens_)));
        final byte[] edk = new byte[encryptResult.getCiphertextBlob().remaining()];
        encryptResult.getCiphertextBlob().get(edk);
        return new DataKey<>(dataKey.getKey(), edk, encryptResult.getKeyId().getBytes(StandardCharsets.UTF_8), this);
    } catch (final AmazonServiceException asex) {
        throw new AwsCryptoException(asex);
    }
}
 

private EncryptResult encrypt0(EncryptRequest req) throws AmazonServiceException, AmazonClientException {
    final byte[] cipherText = new byte[512];
    rnd.nextBytes(cipherText);
    DecryptResult dec = new DecryptResult();
    dec.withKeyId(retrieveArn(req.getKeyId())).withPlaintext(req.getPlaintext().asReadOnlyBuffer());
    ByteBuffer ctBuff = ByteBuffer.wrap(cipherText);

    results_.put(new DecryptMapKey(ctBuff, req.getEncryptionContext()), dec);

    String arn = retrieveArn(req.getKeyId());
    return new EncryptResult().withCiphertextBlob(ctBuff).withKeyId(arn);
}
 

@Test
void testEncryptFails() {
    try {
        textEncryptor.encrypt("Hello");
        failBecauseExceptionWasNotThrown(RuntimeException.class);
    } catch (Exception e) {
        assertThat(e).hasMessageContaining("kmsKeyId must not be blank");
    }

    verify(mockKms, never()).encrypt(any(EncryptRequest.class));
}
 

public static void main(String[] args) {
    AWSKMS kmsClient = AWSKMSClientBuilder.standard().build();

    // Encrypt a data key
    //
    // Replace the following fictitious CMK ARN with a valid CMK ID or ARN
    String keyId = "1234abcd-12ab-34cd-56ef-1234567890ab";
    ByteBuffer plaintext = ByteBuffer.wrap(new byte[]{1, 2, 3, 4, 5, 6, 7, 8, 9, 0});

    EncryptRequest req = new EncryptRequest().withKeyId(keyId).withPlaintext(plaintext);
    ByteBuffer ciphertext = kmsClient.encrypt(req).getCiphertextBlob();
}
 

@Override
public EncryptResult encrypt(EncryptRequest req) throws AmazonServiceException,
        AmazonClientException {
    final byte[] cipherText = new byte[512];
    rnd.nextBytes(cipherText);
    DecryptResult dec = new DecryptResult();
    dec.withKeyId(req.getKeyId()).withPlaintext(req.getPlaintext().asReadOnlyBuffer());
    ByteBuffer ctBuff = ByteBuffer.wrap(cipherText);

    results_.put(new DecryptMapKey(ctBuff, req.getEncryptionContext()), dec);

    return new EncryptResult().withCiphertextBlob(ctBuff).withKeyId(req.getKeyId());
}
 

/**
 * Encrypts the data provided using KMS based on the provided region and key id.
 *
 * @param regionName Region where key is located
 * @param keyId Key id
 * @param data Data to be encrypted
 * @return encrypted data
 */
private byte[] encrypt(final String regionName, final String keyId, final byte[] data) {
  Region region;
  try {
    region = Region.getRegion(Regions.fromName(regionName));
  } catch (IllegalArgumentException iae) {
    throw ApiException.newBuilder()
        .withApiErrors(DefaultApiError.AUTH_IAM_ROLE_AWS_REGION_INVALID)
        .withExceptionCause(iae)
        .build();
  }

  final AWSKMSClient kmsClient = kmsClientFactory.getClient(region);

  try {
    final EncryptResult encryptResult =
        kmsClient.encrypt(
            new EncryptRequest().withKeyId(keyId).withPlaintext(ByteBuffer.wrap(data)));

    return encryptResult.getCiphertextBlob().array();
  } catch (NotFoundException | KMSInvalidStateException keyNotUsableException) {
    throw new KeyInvalidForAuthException(
        String.format("Failed to encrypt token using KMS key with id: %s", keyId),
        keyNotUsableException);
  } catch (AmazonClientException ace) {
    throw ApiException.newBuilder()
        .withApiErrors(DefaultApiError.INTERNAL_SERVER_ERROR)
        .withExceptionCause(ace)
        .withExceptionMessage(
            String.format(
                "Unexpected error communicating with AWS KMS for region %s.", regionName))
        .build();
  }
}
 

@Override
public EncryptResult encrypt(EncryptRequest req) throws AmazonServiceException, AmazonClientException {
    // We internally delegate to encrypt, so as to avoid mockito detecting extra calls to encrypt when spying on the
    // MockKMSClient, we put the real logic into a separate function.
    return encrypt0(req);
}
 

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

    RegionalClientSupplier supplier = mock(RegionalClientSupplier.class);
    when(supplier.getClient(any())).thenReturn(client);

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

    KmsMasterKeyProvider mkp0 = KmsMasterKeyProvider.builder()
                                                    .withDefaultRegion("us-west-2")
                                                    .withCustomClientFactory(supplier)
                                                    .withKeysForEncryption(key1, key2)
                                                    .build();

    MasterKeyProvider<?> mkp = mkp0.withGrantTokens("foo");

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

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

    assertEquals(key1, gdkr.getValue().getKeyId());
    assertEquals(1, gdkr.getValue().getGrantTokens().size());
    assertEquals("foo", gdkr.getValue().getGrantTokens().get(0));

    ArgumentCaptor<EncryptRequest> er = ArgumentCaptor.forClass(EncryptRequest.class);
    verify(client, times(1)).encrypt(er.capture());

    assertEquals(key2, er.getValue().getKeyId());
    assertEquals(1, er.getValue().getGrantTokens().size());
    assertEquals("foo", er.getValue().getGrantTokens().get(0));

    mkp = mkp0.withGrantTokens(Arrays.asList("bar"));

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

    ArgumentCaptor<DecryptRequest> decrypt = ArgumentCaptor.forClass(DecryptRequest.class);
    verify(client, times(1)).decrypt(decrypt.capture());

    assertEquals(1, decrypt.getValue().getGrantTokens().size());
    assertEquals("bar", decrypt.getValue().getGrantTokens().get(0));

    verify(supplier, atLeastOnce()).getClient("us-west-2");
    verifyNoMoreInteractions(supplier);
}
 

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

    RegionalClientSupplier supplier = mock(RegionalClientSupplier.class);
    when(supplier.getClient(any())).thenReturn(client);

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

    KmsMasterKeyProvider mkp0 = KmsMasterKeyProvider.builder()
                                                   .withDefaultRegion("us-west-2")
                                                   .withCustomClientFactory(supplier)
                                                   .withKeysForEncryption(key1, key2)
                                                   .build();
    KmsMasterKey mk1 = mkp0.getMasterKey(key1);
    KmsMasterKey mk2 = mkp0.getMasterKey(key2);

    mk1.setGrantTokens(singletonList("foo"));
    mk2.setGrantTokens(singletonList("foo"));

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

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

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

    assertEquals(key1, gdkr.getValue().getKeyId());
    assertEquals(1, gdkr.getValue().getGrantTokens().size());
    assertEquals("foo", gdkr.getValue().getGrantTokens().get(0));

    ArgumentCaptor<EncryptRequest> er = ArgumentCaptor.forClass(EncryptRequest.class);
    verify(client, times(1)).encrypt(er.capture());

    assertEquals(key2, er.getValue().getKeyId());
    assertEquals(1, er.getValue().getGrantTokens().size());
    assertEquals("foo", er.getValue().getGrantTokens().get(0));

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

    ArgumentCaptor<DecryptRequest> decrypt = ArgumentCaptor.forClass(DecryptRequest.class);
    verify(client, times(1)).decrypt(decrypt.capture());

    assertEquals(1, decrypt.getValue().getGrantTokens().size());
    assertEquals("foo", decrypt.getValue().getGrantTokens().get(0));

    verify(supplier, atLeastOnce()).getClient("us-west-2");
    verifyNoMoreInteractions(supplier);
}
 

private String kmsEncrypt(String value)
{
    String kmsKeyId = context.getSecrets().getSecret("aws.emr.kms_key_id");
    EncryptResult result = kms.encrypt(new EncryptRequest().withKeyId(kmsKeyId).withPlaintext(UTF_8.encode(value)));
    return base64(result.getCiphertextBlob());
}