@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 DataKey<KmsMasterKey> generateDataKey(final CryptoAlgorithm algorithm,
        final Map<String, String> encryptionContext) {
    final GenerateDataKeyResult gdkResult = kms_.get().generateDataKey(updateUserAgent(
            new GenerateDataKeyRequest()
                    .withKeyId(getKeyId())
                    .withNumberOfBytes(algorithm.getDataKeyLength())
                    .withEncryptionContext(encryptionContext)
                    .withGrantTokens(grantTokens_)
    ));
    final byte[] rawKey = new byte[algorithm.getDataKeyLength()];
    gdkResult.getPlaintext().get(rawKey);
    if (gdkResult.getPlaintext().remaining() > 0) {
        throw new IllegalStateException("Recieved an unexpected number of bytes from KMS");
    }
    final byte[] encryptedKey = new byte[gdkResult.getCiphertextBlob().remaining()];
    gdkResult.getCiphertextBlob().get(encryptedKey);

    final SecretKeySpec key = new SecretKeySpec(rawKey, algorithm.getDataKeyAlgo());
    return new DataKey<>(key, encryptedKey, gdkResult.getKeyId().getBytes(StandardCharsets.UTF_8), this);
}
 

@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 generateDataKeyIsCalledWith256NumberOfBits() {
    final AtomicBoolean gdkCalled = new AtomicBoolean(false);
    AWSKMS kmsSpy = new FakeKMS() {
        @Override
        public GenerateDataKeyResult generateDataKey(GenerateDataKeyRequest r) {
            gdkCalled.set(true);
            assertEquals((Integer) 32, r.getNumberOfBytes());
            assertNull(r.getKeySpec());
            return super.generateDataKey(r);
        }
    };
    assertFalse(gdkCalled.get());
    new DirectKmsMaterialProvider(kmsSpy, keyId).getEncryptionMaterials(ctx);
    assertTrue(gdkCalled.get());
}
 

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

}
 

/**
 * @return A key that satisfies the specification defined in BlockCrypto
 */
public EncryptionKey create()
{
    GenerateDataKeyResult dataKeyResult =
            kmsClient.generateDataKey(
                    new GenerateDataKeyRequest()
                            .withKeyId(masterKeyId)
                            .withKeySpec(DataKeySpec.AES_128));

    GenerateRandomRequest randomRequest = new GenerateRandomRequest()
            .withNumberOfBytes(AesGcmBlockCrypto.NONCE_BYTES);
    GenerateRandomResult randomResult = kmsClient.generateRandom(randomRequest);

    return new EncryptionKey(dataKeyResult.getPlaintext().array(), randomResult.getPlaintext().array());
}
 

/**
 * Puts a secret into credstash with a specified version.
 *
 * @param tableName Credstash DynamoDB table name
 * @param secretName Credstash secret name
 * @param secret The secret value
 * @param kmsKeyId The KMS KeyId used to generate a new data key
 * @param context Encryption context for integrity check
 * @param version An optional version string to be used when stashing the secret, defaults to '1' (padded)
 *
 * @throws com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException If the version already exists.
 */
public void putSecret(String tableName, String secretName, String secret, String kmsKeyId, Map<String, String> context, String version) {

    String newVersion = version;
    if(newVersion == null) {
        newVersion = padVersion(1);
    }

    GenerateDataKeyResult generateDataKeyResult = awskmsClient.generateDataKey(new GenerateDataKeyRequest().withKeyId(kmsKeyId).withEncryptionContext(context).withNumberOfBytes(64));
    ByteBuffer plainTextKey = generateDataKeyResult.getPlaintext();
    ByteBuffer cipherTextBlob = generateDataKeyResult.getCiphertextBlob();

    byte[] keyBytes = new byte[32];
    plainTextKey.get(keyBytes);

    byte[] hmacKeyBytes = new byte[plainTextKey.remaining()];
    plainTextKey.get(hmacKeyBytes);

    byte[] encryptedKeyBytes = new byte[cipherTextBlob.remaining()];
    cipherTextBlob.get(encryptedKeyBytes);

    byte[] contents = cryptoImpl.encrypt(keyBytes, secret.getBytes());
    byte[] hmac = cryptoImpl.digest(hmacKeyBytes, contents);

    Map<String, AttributeValue> item = new HashMap<>();
    item.put("name", new AttributeValue(secretName));
    item.put("version", new AttributeValue(newVersion));
    item.put("key", new AttributeValue(new String(Base64.getEncoder().encode(encryptedKeyBytes))));
    item.put("contents", new AttributeValue(new String(Base64.getEncoder().encode(contents))));
    item.put("hmac", new AttributeValue(new String(Hex.encodeHex(hmac))));

    Map<String, String> expressionAttributes = new HashMap<>();
    expressionAttributes.put("#N", "name");

    amazonDynamoDBClient.putItem(new PutItemRequest(tableName, item)
            .withConditionExpression("attribute_not_exists(#N)")
            .withExpressionAttributeNames(expressionAttributes));
}
 

public static void main(String[] args) {
    final String USAGE =
        "To run this example, supply a key id or ARN and a KeySpec\n" +
        "Usage: GenerateDataKey <key-id> <key-spec>\n" +
        "Example: GenerateDataKey 1234abcd-12ab-34cd-56ef-1234567890ab" +
        " AES_256\n";

    if (args.length != 2) {
        System.out.println(USAGE);
        System.exit(1);
    }

    String keyId = args[0];
    String keySpec = args[1];

    AWSKMS kmsClient = AWSKMSClientBuilder.standard().build();

    // Generate a data key

    GenerateDataKeyRequest dataKeyRequest = new GenerateDataKeyRequest();
    dataKeyRequest.setKeyId(keyId);
    dataKeyRequest.setKeySpec(keySpec);

    GenerateDataKeyResult dataKeyResult = kmsClient.generateDataKey(dataKeyRequest);

    ByteBuffer plaintextKey = dataKeyResult.getPlaintext();

    ByteBuffer encryptedKey = dataKeyResult.getCiphertextBlob();

    System.out.printf(
        "Successfully generated an encrypted data key: %s%n",
        Base64.getEncoder().encodeToString(encryptedKey.array())
    );

}
 

@Override
public GenerateDataKeyWithoutPlaintextResult generateDataKeyWithoutPlaintext(
        GenerateDataKeyWithoutPlaintextRequest req) throws AmazonServiceException, AmazonClientException {
    GenerateDataKeyRequest generateDataKeyRequest = new GenerateDataKeyRequest().withEncryptionContext(req.getEncryptionContext())
                                                                                .withGrantTokens(req.getGrantTokens())
                                                                                .withKeyId(req.getKeyId())
                                                                                .withKeySpec(req.getKeySpec())
                                                                                .withNumberOfBytes(req.getNumberOfBytes());
    GenerateDataKeyResult generateDataKey = generateDataKey(generateDataKeyRequest);
    String arn = retrieveArn(req.getKeyId());
    return new GenerateDataKeyWithoutPlaintextResult().withCiphertextBlob(generateDataKey.getCiphertextBlob())
                                                      .withKeyId(arn);
}
 

@Override
public GenerateDataKeyWithoutPlaintextResult generateDataKeyWithoutPlaintext(
        GenerateDataKeyWithoutPlaintextRequest req) throws AmazonServiceException,
        AmazonClientException {
    GenerateDataKeyResult generateDataKey = generateDataKey(new GenerateDataKeyRequest()
            .withEncryptionContext(req.getEncryptionContext()).withNumberOfBytes(
                    req.getNumberOfBytes()));
    return new GenerateDataKeyWithoutPlaintextResult().withCiphertextBlob(
            generateDataKey.getCiphertextBlob()).withKeyId(req.getKeyId());
}
 

protected static String getKMSKey() {
    CreateKeyRequest cmkRequest = new CreateKeyRequest().withDescription("CMK for unit tests");
    CreateKeyResult cmkResult = kmsClient.createKey(cmkRequest);

    GenerateDataKeyRequest dekRequest = new GenerateDataKeyRequest().withKeyId(cmkResult.getKeyMetadata().getKeyId()).withKeySpec("AES_128");
    GenerateDataKeyResult dekResult = kmsClient.generateDataKey(dekRequest);

    return dekResult.getKeyId();
}
 

@Override
public EncryptionMaterials getEncryptionMaterials(EncryptionContext context) {
    final Map<String, String> ec = new HashMap<>();
    ec.put("*" + CONTENT_KEY_ALGORITHM + "*", dataKeyDesc);
    ec.put("*" + SIGNING_KEY_ALGORITHM + "*", sigKeyDesc);
    populateKmsEcFromEc(context, ec);

    final String keyId = selectEncryptionKeyId(context);
    if (StringUtils.isNullOrEmpty(keyId)) {
        throw new DynamoDBMappingException("Encryption key id is empty.");
    }

    final GenerateDataKeyRequest req = appendUserAgent(new GenerateDataKeyRequest());
    req.setKeyId(keyId);
    // NumberOfBytes parameter is used because we're not using this key as an AES-256 key,
    // we're using it as an HKDF-SHA256 key.
    req.setNumberOfBytes(256 / 8);
    req.setEncryptionContext(ec);

    final GenerateDataKeyResult dataKeyResult = generateDataKey(req, context);

    final Map<String, String> materialDescription = new HashMap<>();
    materialDescription.putAll(description);
    materialDescription.put(COVERED_ATTR_CTX_KEY, KEY_COVERAGE);
    materialDescription.put(KEY_WRAPPING_ALGORITHM, "kms");
    materialDescription.put(CONTENT_KEY_ALGORITHM, dataKeyDesc);
    materialDescription.put(SIGNING_KEY_ALGORITHM, sigKeyDesc);
    materialDescription.put(ENVELOPE_KEY, Base64.encodeToString(toArray(dataKeyResult.getCiphertextBlob())));

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

    kdf.init(toArray(dataKeyResult.getPlaintext()));

    final SecretKey encryptionKey = new SecretKeySpec(kdf.deriveKey(KDF_ENC_INFO, dataKeyLength / 8), dataKeyAlg);
    final SecretKey signatureKey = new SecretKeySpec(kdf.deriveKey(KDF_SIG_INFO, sigKeyLength / 8), sigKeyAlg);
    return new SymmetricRawMaterials(encryptionKey, signatureKey, materialDescription);
}
 

@Override
protected GenerateDataKeyResult generateDataKey(GenerateDataKeyRequest request, EncryptionContext context) {
    return super.generateDataKey(request, context);
}
 

/**
 * Returns a data encryption key that you can use in your application to encrypt data locally. The default
 * implementation calls KMS to generate the data key using the parameters provided in the
 * {@link GenerateDataKeyRequest}. Subclass can override the default implementation to provide additional
 * request parameters using attributes within the {@link EncryptionContext}.
 *
 * @param request request parameters to generate the data key.
 * @param context additional useful data to generate the data key.
 * @return the newly generated data key which includes both the plaintext and ciphertext.
 */
protected GenerateDataKeyResult generateDataKey(final GenerateDataKeyRequest request,
        final EncryptionContext context) {
    return kms.generateDataKey(request);
}