Java 类com.amazonaws.services.kms.model.DecryptRequest 实例源码

@Before
public void setUp() {
    dynamoDBClient = Mockito.mock(AmazonDynamoDB.class);

    GenerateDataKeyResult generateDatakeyResult = new GenerateDataKeyResult();
    generateDatakeyResult.setCiphertextBlob(Mockito.mock(ByteBuffer.class));
    generateDatakeyResult.setPlaintext(Mockito.mock(ByteBuffer.class));

    DecryptResult decryptResult = new DecryptResult();
    decryptResult.setKeyId("alias/foo");
    decryptResult.setPlaintext(Mockito.mock(ByteBuffer.class));

    awskmsClient = Mockito.mock(AWSKMS.class);
    Mockito.when(awskmsClient.generateDataKey(Mockito.any(GenerateDataKeyRequest.class))).thenReturn(generateDatakeyResult);
    Mockito.when(awskmsClient.decrypt(Mockito.any(DecryptRequest.class))).thenReturn(decryptResult);
}

@Override
public DecryptResult decrypt(DecryptRequest req) throws AmazonServiceException, AmazonClientException {
    DecryptResult result = results_.get(new DecryptMapKey(req));
    if (result != null) {
        // Copy it to avoid external modification
        DecryptResult copy = new DecryptResult();
        copy.setKeyId(retrieveArn(result.getKeyId()));
        byte[] pt = new byte[result.getPlaintext().limit()];
        result.getPlaintext().get(pt);
        result.getPlaintext().rewind();
        copy.setPlaintext(ByteBuffer.wrap(pt));
        return copy;
    } else {
        throw new InvalidCiphertextException("Invalid Ciphertext");
    }
}

public static String decrypt(String str, Region region) throws UnsupportedEncodingException {
  if (isJUnitTest()) {
    return str;
  }

  AWSKMS kms = AWSKMSClientBuilder.standard().withRegion(region.getName()).build();

  /*
   * The KMS ciphertext is base64 encoded and must be decoded before the request is made
   */
  String cipherString = str;
  byte[] cipherBytes = Base64.decode(cipherString);

  /*
   * Create decode request and decode
   */
  ByteBuffer cipherBuffer = ByteBuffer.wrap(cipherBytes);
  DecryptRequest req = new DecryptRequest().withCiphertextBlob(cipherBuffer);
  DecryptResult resp = kms.decrypt(req);

  /*
   * Convert the response plaintext bytes to a string
   */
  return new String(resp.getPlaintext().array(), Charset.forName("UTF-8"));
}

/**
 * @param dataIn
 * @return returns byte[]
 * @throws FaultResponse
 */
@Override
public String accessGenericData(String dataIn) throws FaultResponse {
    byte[] decodedData = AES.decode(dataIn);
    byte[] encryptedKeyBytes = Arrays.copyOfRange(decodedData,0, 151);
    byte[] originalEncrypted = Arrays.copyOfRange(decodedData, 151, decodedData.length);

    // Decrypt the encrypted key to get the original plain text key...
    ByteBuffer ciphertextBlob = ByteBuffer.wrap(encryptedKeyBytes);

    DecryptRequest req = new DecryptRequest().withCiphertextBlob(ciphertextBlob);
    ByteBuffer plainText = kms.decrypt(req).getPlaintext();

    String ret;
    try {
        ret = new String(AES.unPackBytes(AES.decrypt(originalEncrypted, plainText.array())));
    } catch (Exception e) {
        e.printStackTrace();
        Fault fault = new Fault();
        fault.setErrorCode(500);
        throw new FaultResponse("Error doing AES decryption: " + e.getLocalizedMessage(), fault);
    }

    return ret;
}

@Override
public KeyProvider build() {
    if ( null == key || 0 == key.length ) {
        return new KeyProviderImpl(null);
    } else if ( 16 == key.length ) {
        return new KeyProviderImpl(new SecretKeySpec(key, "AES"));
    }
    AWSKMS kms = _amazonWebServiceClients.withEndpoint(
        new AWSKMSClient(
            _credProviderFactory.create(credProvider),
            _clientConfigurations.withProxy(new ClientConfiguration(), proxy)),
        endpoint);
    key = kms.decrypt(new DecryptRequest()
                      .withCiphertextBlob(ByteBuffer.wrap(key)))
        .getPlaintext().array();
    if ( 16 != key.length ) {
        LOG.warn("Expected decrypted key to be exactly 16 bytes, got "+key.length+" bytes. Please "+
                 "verify the key was not base64 encoded before encrypting with KMS");
        return new KeyProviderImpl(null);
    }
    return new KeyProviderImpl(new SecretKeySpec(key, "AES"));
}

/**
 * Decodes the encrypted token and attempts to decrypt it using AWS KMS. If
 * successful, the token is returned.
 *
 * @param kmsClient      KMS client
 * @param encryptedToken Token to decode and decrypt
 * @return Decrypted token
 */
protected VaultAuthResponse decryptToken(AWSKMS kmsClient, String encryptedToken) {
    byte[] decodedToken;

    try {
        decodedToken = Base64.decode(encryptedToken);
    } catch (IllegalArgumentException iae) {
        throw new VaultClientException("Encrypted token not Base64 encoded", iae);
    }

    final DecryptRequest request = new DecryptRequest().withCiphertextBlob(ByteBuffer.wrap(decodedToken));
    final DecryptResult result = kmsClient.decrypt(request);

    final String decryptedAuthData = new String(result.getPlaintext().array(), Charset.forName("UTF-8"));

    return gson.fromJson(decryptedAuthData, VaultAuthResponse.class);
}

@Test
public void testAwsPrivateKeyStore() throws Exception {
    String bucketName = "my_bucket";
    String keyName = "my_key";
    String expected = "my_value";

    AmazonS3 s3 = Mockito.mock(AmazonS3.class);
    AWSKMS kms = Mockito.mock(AWSKMS.class);
    S3Object s3Object = Mockito.mock(S3Object.class);
    Mockito.when(s3.getObject(bucketName, keyName)).thenReturn(s3Object);
    InputStream is = new ByteArrayInputStream( expected.getBytes() );
    S3ObjectInputStream s3ObjectInputStream = new S3ObjectInputStream(is, null);
    Mockito.when(s3Object.getObjectContent()).thenReturn(s3ObjectInputStream);

    String result = expected;
    ByteBuffer buffer = ByteBuffer.wrap(result.getBytes());
    DecryptResult decryptResult = Mockito.mock(DecryptResult.class); 
    Mockito.when(kms.decrypt(Mockito.any(DecryptRequest.class))).thenReturn(decryptResult);
    Mockito.when(decryptResult.getPlaintext()).thenReturn(buffer);

    AwsPrivateKeyStore awsPrivateKeyStore = new AwsPrivateKeyStore(s3, kms);
    String actual = awsPrivateKeyStore.getApplicationSecret(bucketName, keyName);
    Assert.assertEquals(actual, expected);

}

@Override
public String decrypt(final String encryptedText) {
    if (encryptedText == null || encryptedText.isEmpty()) {
        return EMPTY_STRING;
    } else {

        // Extract the encryption context and the remaining part
        final Map<String, String> encryptionContext = extractEncryptionContext(encryptedText);
        final String encryptedValue = extractEncryptedValue(encryptedText);

        // Assuming the encryptedText is encoded in Base64
        final ByteBuffer encryptedBytes =
                ByteBuffer.wrap(Base64.decode(encryptedValue.getBytes()));

        final DecryptRequest decryptRequest = new DecryptRequest()
                .withCiphertextBlob(encryptedBytes)
                .withEncryptionContext(encryptionContext);

        return extractString(kms.decrypt(decryptRequest).getPlaintext());
    }
}

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

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

    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()));

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

@Override
public String decrypt(AwsParamsDto awsParamsDto, String base64ciphertextBlob)
{
    // Construct a new AWS KMS service client using the specified client configuration.
    // A credentials provider chain will be used that searches for credentials in this order:
    // - Environment Variables - AWS_ACCESS_KEY_ID and AWS_SECRET_KEY
    // - Java System Properties - aws.accessKeyId and aws.secretKey
    // - Instance Profile Credentials - delivered through the Amazon EC2 metadata service
    AWSKMSClient awsKmsClient = new AWSKMSClient(awsHelper.getClientConfiguration(awsParamsDto));

    // Decode the base64 encoded ciphertext.
    ByteBuffer ciphertextBlob = ByteBuffer.wrap(Base64.decodeBase64(base64ciphertextBlob));

    // Create the decrypt request.
    DecryptRequest decryptRequest = new DecryptRequest().withCiphertextBlob(ciphertextBlob);

    // Call AWS KMS decrypt service method.
    DecryptResult decryptResult = kmsOperations.decrypt(awsKmsClient, decryptRequest);

    // Get decrypted plaintext data.
    ByteBuffer plainText = decryptResult.getPlaintext();

    // Return the plain text as a string.
    return new String(plainText.array(), StandardCharsets.UTF_8);
}

@Override
public DecryptResult decrypt(AWSKMSClient awsKmsClient, DecryptRequest decryptRequest)
{
    // Check the cipher text.
    if (decryptRequest.getCiphertextBlob().equals(ByteBuffer.wrap(Base64.decodeBase64(MOCK_CIPHER_TEXT_INVALID))))
    {
        throw new InvalidCiphertextException("(Service: AWSKMS; Status Code: 400; Error Code: InvalidCiphertextException; Request ID: NONE)");
    }

    DecryptResult decryptResult = new DecryptResult();

    // Convert the test plain text to byte buffer and set the plain text return value.
    decryptResult.setPlaintext(ByteBuffer.wrap(MOCK_PLAIN_TEXT.getBytes()));

    return decryptResult;
}

@Test
public void testEncryptDecrypt() throws Exception {
  DecryptResult mockDecryptResult = mock(DecryptResult.class);
  EncryptResult mockEncryptResult = mock(EncryptResult.class);
  when(mockKms.decrypt(isA(DecryptRequest.class)))
      .thenReturn(mockDecryptResult);
  when(mockKms.encrypt(isA(EncryptRequest.class)))
      .thenReturn(mockEncryptResult);

  Aead aead = new AwsKmsAead(mockKms, keyId);
  byte[] aad = Random.randBytes(20);
  for (int messageSize = 0; messageSize < 75; messageSize++) {
    byte[] message = Random.randBytes(messageSize);
    when(mockDecryptResult.getPlaintext()).thenReturn(ByteBuffer.wrap(message));
    when(mockEncryptResult.getCiphertextBlob()).thenReturn(ByteBuffer.wrap(message));
    byte[] ciphertext = aead.encrypt(message, aad);
    byte[] decrypted = aead.decrypt(ciphertext, aad);
    assertArrayEquals(message, decrypted);
  }
}

@Test
public void testDecrypt_shouldThrowExceptionIfRequestFailed() throws Exception {
  EncryptResult mockEncryptResult = mock(EncryptResult.class);
  when(mockKms.encrypt(isA(EncryptRequest.class)))
      .thenReturn(mockEncryptResult);
  AmazonServiceException exception = mock(AmazonServiceException.class);
  when(mockKms.decrypt(isA(DecryptRequest.class)))
      .thenThrow(exception);

  Aead aead = new AwsKmsAead(mockKms, keyId);
  byte[] aad = Random.randBytes(20);
  byte[] message = Random.randBytes(20);
  when(mockEncryptResult.getCiphertextBlob()).thenReturn(ByteBuffer.wrap(message));
  byte[] ciphertext = aead.encrypt(message, aad);
  try {
    aead.decrypt(ciphertext, aad);
    fail("Expected GeneralSecurityException");
  } catch (GeneralSecurityException e) {
    // expected.
  }
}

/**
 * Decrypts the secured CEK via KMS; involves network calls.
 *
 * @return the CEK (in plaintext).
 */
private static SecretKey cekByKMS(byte[] cekSecured, String keyWrapAlgo,
        EncryptionMaterials materials,
        ContentCryptoScheme contentCryptoScheme, AWSKMS kms) {
    DecryptRequest kmsreq = new DecryptRequest()
        .withEncryptionContext(materials.getMaterialsDescription())
        .withCiphertextBlob(ByteBuffer.wrap(cekSecured));
    DecryptResult result = kms.decrypt(kmsreq);
    return new SecretKeySpec(copyAllBytesFrom(result.getPlaintext()),
            contentCryptoScheme.getKeyGeneratorAlgorithm());
}

private String decryptString(byte[] encryptedString) {
    DecryptRequest request = new DecryptRequest();
    LOGGER.fine("decrypting with kms");
    if (null != this.kmsEncryptionContextKey && null != this.kmsSecretName) {
        LOGGER.info("decrypting with context");
        request.addEncryptionContextEntry(this.kmsEncryptionContextKey, this.kmsSecretName);
    }
    request.setCiphertextBlob(ByteBuffer.wrap(encryptedString));
    DecryptResult decryptResult = this.amazonKmsClientBuilder.build().decrypt(request);
    LOGGER.fine("decrypted with kms");
    return charset.decode(decryptResult.getPlaintext()).toString();
}

@Test
public void passwordUsesTheS3Bucket() throws Exception {
    S3Object mockS3Object = mock(S3Object.class);
    AmazonS3Client mockClient = mock(AmazonS3Client.class);
    when(mockClientBuilder.build()).thenReturn(mockClient);
    when(mockClient.getObject(any(GetObjectRequest.class))).thenReturn(mockS3Object);
    AWSKMSClient mockKmsClient = mock(AWSKMSClient.class);
    when(mockKmsClientBuilder.build()).thenReturn(mockKmsClient);

    S3ObjectInputStream mockS3ObjectInputStream = mock(S3ObjectInputStream.class);
    when(mockS3Object.getObjectContent()).thenReturn(mockS3ObjectInputStream);
    when(mockS3ObjectInputStream.read(new byte[anyInt()], anyInt(), anyByte()))
            .thenAnswer(new WriteBufferAnswer("encryptedPassword".getBytes()))
            .thenReturn(-1);

    DecryptResult result = new DecryptResult();
    CharsetEncoder charsetEncoder = Charset.forName("UTF-8").newEncoder();
    result.setPlaintext(charsetEncoder.encode(CharBuffer.wrap("password")));
    when(mockKmsClient.decrypt(any(DecryptRequest.class))).thenReturn(result);

    Secret secret = test.getPassword();

    // have we got the expected password
    assertThat(secret.getPlainText()).isEqualTo("password");

    // have we used the bucket
    ArgumentCaptor<GetObjectRequest> capturedObjectRequest = ArgumentCaptor.forClass(GetObjectRequest.class);
    verify(mockClient).getObject(capturedObjectRequest.capture());
    assertThat(capturedObjectRequest.getValue().getBucketName()).isEqualTo("bucketUri");
    assertThat(capturedObjectRequest.getValue().getS3ObjectId().getKey()).isEqualTo("/bucketPath");

    // have we used kms to decrypt
    ArgumentCaptor<DecryptRequest> capturedDecryptRequest = ArgumentCaptor.forClass(DecryptRequest.class);
    verify(mockKmsClient).decrypt(capturedDecryptRequest.capture());
    assertThat(capturedDecryptRequest.getValue().getEncryptionContext()).containsEntry("someEncryptContextKey", "kmsEncryptContextValue");
    ByteBuffer ciphertextBlob = capturedDecryptRequest.getValue().getCiphertextBlob();
    assertThat(new String(Charset.forName("UTF-8").decode(ciphertextBlob).array())).isEqualTo("encryptedPassword");

}

@Test
public void closesIfIOExceptionWhileReading() throws Exception {
    AmazonS3Client mockClient = mock(AmazonS3Client.class);
    when(mockClientBuilder.build()).thenReturn(mockClient);

    AWSKMSClient mockKmsClient = mock(AWSKMSClient.class);
    when(mockKmsClientBuilder.build()).thenReturn(mockKmsClient);

    S3Object mockS3Object = mock(S3Object.class);
    when(mockClient.getObject(any(GetObjectRequest.class))).thenReturn(mockS3Object);
    S3ObjectInputStream mockS3ObjectInputStream = mock(S3ObjectInputStream.class);
    when(mockS3Object.getObjectContent()).thenReturn(mockS3ObjectInputStream);
    when(mockS3ObjectInputStream.read(new byte[anyInt()], anyInt(), anyByte()))
            .thenAnswer(new WriteBufferAnswer("encryptedPassword".getBytes()))
            .thenThrow(new IOException("something went wrong"))
            .thenReturn(-1);

    DecryptResult result = new DecryptResult();
    CharsetEncoder charsetEncoder = Charset.forName("UTF-8").newEncoder();
    result.setPlaintext(charsetEncoder.encode(CharBuffer.wrap("password")));
    when(mockKmsClient.decrypt(any(DecryptRequest.class))).thenReturn(result);

    Secret secret = null;
    try {
        secret = test.getPassword();
        TestCase.fail("should have thrown exception");
    } catch (AwsBucketReadingException e) {
        assertThat(e.getCause()).isInstanceOf(IOException.class);
    }

    // have we used the bucket
    ArgumentCaptor<GetObjectRequest> capturedObjectRequest = ArgumentCaptor.forClass(GetObjectRequest.class);
    verify(mockClient).getObject(capturedObjectRequest.capture());
    assertThat(capturedObjectRequest.getValue().getBucketName()).isEqualTo("bucketUri");
    assertThat(capturedObjectRequest.getValue().getS3ObjectId().getKey()).isEqualTo("/bucketPath");

    // and we have closed it even if there was an exception
    verify(mockS3Object).close();
}

@Override
public DataKey<KmsMasterKey> decryptDataKey(final CryptoAlgorithm algorithm,
        final Collection<? extends EncryptedDataKey> encryptedDataKeys,
        final Map<String, String> encryptionContext)
        throws UnsupportedProviderException, AwsCryptoException {
    final List<Exception> exceptions = new ArrayList<>();
    for (final EncryptedDataKey edk : encryptedDataKeys) {
        try {
            final DecryptResult decryptResult = kms_.decrypt(
                    new DecryptRequest()
                            .withCiphertextBlob(ByteBuffer.wrap(edk.getEncryptedDataKey()))
                            .withEncryptionContext(encryptionContext)
                            .withGrantTokens(grantTokens_));
            if (decryptResult.getKeyId().equals(id_)) {
                final byte[] rawKey = new byte[algorithm.getDataKeyLength()];
                decryptResult.getPlaintext().get(rawKey);
                if (decryptResult.getPlaintext().remaining() > 0) {
                    throw new IllegalStateException("Received an unexpected number of bytes from KMS");
                }
                return new DataKey<>(
                        new SecretKeySpec(rawKey, algorithm.getDataKeyAlgo()),
                        edk.getEncryptedDataKey(),
                        edk.getProviderInformation(), this);
            }
        } catch (final AmazonServiceException awsex) {
            exceptions.add(awsex);
        }
    }

    throw buildCannotDecryptDksException(exceptions);
}

public DecryptMapKey(DecryptRequest req) {
    cipherText = req.getCiphertextBlob().asReadOnlyBuffer();
    if (req.getEncryptionContext() != null) {
        ec = Collections.unmodifiableMap(new HashMap<String, String>(req.getEncryptionContext()));
    } else {
        ec = Collections.emptyMap();
    }
}

/**
 * A utility function to decode and decrypt KMS encrypted data for testing
 * @param encryptedData The encoded, encrypted data to encrypt
 * @param encryptor An initialized KMS encryptor
 * @return The decrpted data as a byte array
 */
public static byte[] kmsDecrypt(String encryptedData, KMSEncryptor encryptor) throws Exception {
    byte[] decodedData = AES.decode(encryptedData);
    byte[] encryptedKeyBytes = Arrays.copyOfRange(decodedData,0, 151);
    byte[] originalEncrypted = Arrays.copyOfRange(decodedData, 151, decodedData.length);

    // Decrypt the encrypted key to get the original plain text key...
    ByteBuffer ciphertextBlob = ByteBuffer.wrap(encryptedKeyBytes);

    DecryptRequest req = new DecryptRequest().withCiphertextBlob(ciphertextBlob);
    ByteBuffer plainText = encryptor.kms.decrypt(req).getPlaintext();

    return AES.decrypt(originalEncrypted, plainText.array());
}

/**
 * A utility function to decode and decrypt KMS encrypted data for testing
 * @param encryptedData The encoded, encrypted data to encrypt
 * @param encryptor An initialized KMS encryptor
 * @return The decrpted data as a byte array
 */
public static byte[] kmsDecrypt(String encryptedData, KMSEncryptor encryptor) throws Exception {
    byte[] decodedData = AES.decode(encryptedData);
    byte[] encryptedKeyBytes = Arrays.copyOfRange(decodedData,0, 151);
    byte[] originalEncrypted = Arrays.copyOfRange(decodedData, 151, decodedData.length);

    // Decrypt the encrypted key to get the original plain text key...
    ByteBuffer ciphertextBlob = ByteBuffer.wrap(encryptedKeyBytes);

    DecryptRequest req = new DecryptRequest().withCiphertextBlob(ciphertextBlob);
    ByteBuffer plainText = encryptor.kms.decrypt(req).getPlaintext();

    return AES.decrypt(originalEncrypted, plainText.array());
}

@Override
public Observable<byte[]> decrypt(VertxContext<Server> vertxContext, byte[] cipherBytes) {
    SfsVertx sfsVertx = vertxContext.vertx();
    Context context = sfsVertx.getOrCreateContext();
    return Observable.defer(() -> RxHelper.executeBlocking(context, sfsVertx.getBackgroundPool(), () -> {
        DecryptRequest req =
                new DecryptRequest()
                        .withCiphertextBlob(ByteBuffer.wrap(cipherBytes.clone()));
        ByteBuffer buffer = kms.decrypt(req).getPlaintext();
        byte[] b = new byte[buffer.remaining()];
        buffer.get(b);
        return b;
    }));
}

String getDecryptedData(final String bucketName, final String keyName) {

    String keyValue = "";
    S3Object s3Object = s3.getObject(bucketName, keyName);

    if (LOG.isDebugEnabled()) {
        LOG.debug("retrieving appName {}, key {}", bucketName, keyName);
    }

    if (null == s3Object) {
        LOG.error("error retrieving key {}, from bucket {}", keyName, bucketName);
        return keyValue;
    }

    try (S3ObjectInputStream s3InputStream = s3Object.getObjectContent(); 
            ByteArrayOutputStream result = new ByteArrayOutputStream();) {

        byte[] buffer = new byte[1024];
        int length;
        while ((length = s3InputStream.read(buffer)) != -1) {
            result.write(buffer, 0, length);
        }

        // if key should be decrypted, do so with KMS

        if (kmsDecrypt) {
            DecryptRequest req = new DecryptRequest().withCiphertextBlob(ByteBuffer.wrap(result.toByteArray()));
            ByteBuffer plainText = kms.decrypt(req).getPlaintext();
            keyValue = new String(plainText.array());
        } else {
            keyValue = result.toString();
        }

    } catch (IOException e) {
        LOG.error("error getting application secret.", e);
    }

    return keyValue.trim();
}

/**
 * Decrypts the given encrypted bytes with the given encryption context.
 * Assumes the runtime has access to the desired AWS environment.
 *
 * @param region            the region
 * @param encryptionContext optional encryption context
 * @param encryptedBytes    the encrypted bytes
 * @return the decrypted bytes
 */
public static byte[] decrypt(Region region, Map<String, String> encryptionContext, byte[] encryptedBytes) {
    DecryptRequest decryptRequest = new DecryptRequest();
    decryptRequest.setEncryptionContext(encryptionContext);
    decryptRequest.setCiphertextBlob(ByteBuffer.wrap(encryptedBytes));

    AWSKMSClient client = new AWSKMSClient();
    client.setRegion(region);
    DecryptResult decryptResult = client.decrypt(decryptRequest);

    return decryptResult.getPlaintext().array();
}

@Bean
AWSKMS kms() {
    final AWSKMS mock = mock(AWSKMS.class);
    when(mock.decrypt(any(DecryptRequest.class))).thenAnswer((Answer<DecryptResult>) invocation ->
            new DecryptResult().withPlaintext(ByteBuffer.wrap(KmsEncryptionIntegrationConfigurationTest.PLAINTEXT.getBytes())));
    return mock;
}

@Test
public void testPropertyHasBeenDecrypted() throws Exception {

    assertThat(decryptedSecret).isEqualTo(PLAINTEXT);

    final DecryptRequest decryptRequest = new DecryptRequest();
    decryptRequest.setCiphertextBlob(CIPHER_TEXT_BLOB);
    verify(mockKms, atLeastOnce()).decrypt(decryptRequest);
}

@Override
public DecryptionMaterials getDecryptionMaterials(EncryptionContext context) {
    final Map<String, String> materialDescription = context.getMaterialDescription();

    final Map<String, String> ec = new HashMap<>();
    final String providedEncAlg = materialDescription.get(CONTENT_KEY_ALGORITHM);
    final String providedSigAlg = materialDescription.get(SIGNING_KEY_ALGORITHM);

    ec.put("*" + CONTENT_KEY_ALGORITHM + "*", providedEncAlg);
    ec.put("*" + SIGNING_KEY_ALGORITHM + "*", providedSigAlg);

    populateKmsEcFromEc(context, ec);

    DecryptRequest request = appendUserAgent(new DecryptRequest());
    request.setCiphertextBlob(ByteBuffer.wrap(Base64.decode(materialDescription.get(ENVELOPE_KEY))));
    request.setEncryptionContext(ec);
    final DecryptResult decryptResult = kms.decrypt(request);
    validateEncryptionKeyId(decryptResult.getKeyId(), context);

    final Hkdf kdf;
    try {
        kdf = Hkdf.getInstance(KDF_ALG);
    } catch (NoSuchAlgorithmException e) {
        throw new DynamoDBMappingException(e);
    }
    kdf.init(toArray(decryptResult.getPlaintext()));

    final String[] encAlgParts = providedEncAlg.split("/", 2);
    int encLength = encAlgParts.length == 2 ? Integer.parseInt(encAlgParts[1]) : 256;
    final String[] sigAlgParts = providedSigAlg.split("/", 2);
    int sigLength = sigAlgParts.length == 2 ? Integer.parseInt(sigAlgParts[1]) : 256;

    final SecretKey encryptionKey = new SecretKeySpec(kdf.deriveKey(KDF_ENC_INFO, encLength / 8), encAlgParts[0]);
    final SecretKey macKey = new SecretKeySpec(kdf.deriveKey(KDF_SIG_INFO, sigLength / 8), sigAlgParts[0]);

    return new SymmetricRawMaterials(encryptionKey, macKey, materialDescription);
}

@Override
public DecryptResult decrypt(DecryptRequest req) throws AmazonServiceException,
        AmazonClientException {
    DecryptResult result = results_.get(new DecryptMapKey(req));
    if (result != null) {
        return result;
    } else {
        throw new InvalidCiphertextException("Invalid Ciphertext");
    }
}

public DecryptMapKey(DecryptRequest req) {
    cipherText = req.getCiphertextBlob().asReadOnlyBuffer();
    if (req.getEncryptionContext() != null) {
        ec = Collections.unmodifiableMap(new HashMap<String, String>(req
                .getEncryptionContext()));
    } else {
        ec = Collections.emptyMap();
    }
}

@Override
public byte[] decrypt(final byte[] ciphertext, final byte[] associatedData)
    throws GeneralSecurityException {
  try {
    DecryptRequest req = new DecryptRequest().withCiphertextBlob(ByteBuffer.wrap(ciphertext));
    if (associatedData != null && associatedData.length != 0) {
      req = req.addEncryptionContextEntry("associatedData", BinaryUtils.toHex(associatedData));
    }
    return kmsClient.decrypt(req).getPlaintext().array();
  } catch (AmazonServiceException e) {
    throw new GeneralSecurityException("decryption failed", e);
  }
}

@Test
public void doesNotUseEncryptContextIfNotProvided() throws Exception {
    AwsBucketCredentialsImpl test = new AwsBucketCredentialsImpl(CredentialsScope.GLOBAL, "myId",
            "EU_WEST_1", "bucketUri", "/bucketPath", "username", true,
            "mydescription", null, "kmsEncryptContextValue", true, "host", "9000");
    Whitebox.setInternalState(test, "amazonS3ClientBuilder", mockClientBuilder);
    Whitebox.setInternalState(test, "amazonKmsClientBuilder", mockKmsClientBuilder);

    AmazonS3Client mockClient = mock(AmazonS3Client.class);
    when(mockClientBuilder.build()).thenReturn(mockClient);

    AWSKMSClient mockKmsClient = mock(AWSKMSClient.class);
    when(mockKmsClientBuilder.build()).thenReturn(mockKmsClient);
    S3Object mockS3Object = mock(S3Object.class);
    when(mockClient.getObject(any(GetObjectRequest.class))).thenReturn(mockS3Object);
    S3ObjectInputStream mockS3ObjectInputStream = mock(S3ObjectInputStream.class);
    when(mockS3Object.getObjectContent()).thenReturn(mockS3ObjectInputStream);
    when(mockS3ObjectInputStream.read(new byte[anyInt()], anyInt(), anyByte()))
            .thenAnswer(new WriteBufferAnswer("encryptedPassword".getBytes()))
            .thenReturn(-1);

    DecryptResult result = new DecryptResult();
    CharsetEncoder charsetEncoder = Charset.forName("UTF-8").newEncoder();
    result.setPlaintext(charsetEncoder.encode(CharBuffer.wrap("password")));
    when(mockKmsClient.decrypt(any(DecryptRequest.class))).thenReturn(result);

    Secret secret = test.getPassword();

    // have we got the expected password
    assertThat(secret.getPlainText()).isEqualTo("password");

    // have we used the bucket
    ArgumentCaptor<GetObjectRequest> capturedObjectRequest = ArgumentCaptor.forClass(GetObjectRequest.class);
    verify(mockClient).getObject(capturedObjectRequest.capture());
    assertThat(capturedObjectRequest.getValue().getBucketName()).isEqualTo("bucketUri");
    assertThat(capturedObjectRequest.getValue().getS3ObjectId().getKey()).isEqualTo("/bucketPath");

    // have we used kms to decrypt
    ArgumentCaptor<DecryptRequest> capturedDecryptRequest = ArgumentCaptor.forClass(DecryptRequest.class);
    verify(mockKmsClient).decrypt(capturedDecryptRequest.capture());
    assertThat(capturedDecryptRequest.getValue().getEncryptionContext()).isEmpty();
    ByteBuffer ciphertextBlob = capturedDecryptRequest.getValue().getCiphertextBlob();
    assertThat(new String(Charset.forName("UTF-8").decode(ciphertextBlob).array())).isEqualTo("encryptedPassword");

    verify(mockS3Object).close();

}

private ByteBuffer decrypt(ByteBuffer ciphertextblob) {

        DecryptRequest req = new DecryptRequest().withCiphertextBlob(ciphertextblob);
        return kms.decrypt(req).getPlaintext();
    }

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

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

protected void mockDecrypt(AWSKMSClient kmsClient, final String toDecrypt) {
    DecryptResult decryptResult = new DecryptResult();
    decryptResult.setPlaintext(ByteBuffer.wrap(toDecrypt.getBytes()));
    when(kmsClient.decrypt(any(DecryptRequest.class))).thenReturn(decryptResult);
}

@Override
public DecryptResult decrypt(AWSKMSClient awsKmsClient, DecryptRequest decryptRequest)
{
    // Call AWS KMS decrypt service method.
    return awsKmsClient.decrypt(decryptRequest);
}

private ByteBuffer decryptValue(final byte[] encryptedBytes) {
    final DecryptRequest request = new DecryptRequest().withCiphertextBlob(ByteBuffer.wrap(encryptedBytes));
    return Failsafe.with(retryPolicy).get(() -> kmsClient.decrypt(request).getPlaintext());
}

protected ByteBuffer decryptKeyWithKMS(byte[] encryptedKeyBytes, Map<String, String> context) {
    ByteBuffer blob = ByteBuffer.wrap(encryptedKeyBytes);

    DecryptResult decryptResult = awskmsClient.decrypt(new DecryptRequest().withCiphertextBlob(blob).withEncryptionContext(context));

    return decryptResult.getPlaintext();
}

/**
 * Executes the decrypt request by calling the AWS KMS service.
 *
 * @param awsKmsClient the client for accessing the AWS KMS service
 * @param decryptRequest the decrypt request
 *
 * @return the response from the decrypt service method, as returned by AWS KMS service
 */
public DecryptResult decrypt(AWSKMSClient awsKmsClient, DecryptRequest decryptRequest);