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

@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_.generateDataKey(
            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);
}

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

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

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

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

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

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

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

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

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

}

/**
 * @param materials a non-null encryption material
 */
private ContentCryptoMaterial buildContentCryptoMaterial(
        EncryptionMaterials materials, Provider provider,
        AmazonWebServiceRequest req) {
    // Randomly generate the IV
    final byte[] iv = new byte[contentCryptoScheme.getIVLengthInBytes()];
    cryptoScheme.getSecureRandom().nextBytes(iv);

    if (materials.isKMSEnabled()) {
        final Map<String, String> encryptionContext =
                ContentCryptoMaterial.mergeMaterialDescriptions(materials, req);
        GenerateDataKeyRequest keyGenReq = new GenerateDataKeyRequest()
            .withEncryptionContext(encryptionContext)
            .withKeyId(materials.getCustomerMasterKeyId())
            .withKeySpec(contentCryptoScheme.getKeySpec());
        keyGenReq
            .withGeneralProgressListener(req.getGeneralProgressListener())
            .withRequestMetricCollector(req.getRequestMetricCollector())
            ;
        GenerateDataKeyResult keyGenRes = kms.generateDataKey(keyGenReq);
        final SecretKey cek =
            new SecretKeySpec(copyAllBytesFrom(keyGenRes.getPlaintext()),
                    contentCryptoScheme.getKeyGeneratorAlgorithm());
        byte[] keyBlob = copyAllBytesFrom(keyGenRes.getCiphertextBlob());
        return ContentCryptoMaterial.wrap(cek, iv,
                contentCryptoScheme, provider,
                new KMSSecuredCEK(keyBlob, encryptionContext));
    } else {
        // Generate a one-time use symmetric key and initialize a cipher to encrypt object data
        return ContentCryptoMaterial.create(
                generateCEK(materials, provider),
                iv, materials, cryptoScheme, provider, kms, req);
    }
}

/**
 * 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));
}

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

/**
 * @param dataIn a String representing a value that has to be encrypted.
 * @return returns byte[]
 * @throws FaultResponse if the encoding can't be done.
 */
@Override
public String protectGenericData(String dataIn) throws FaultResponse {
    //Get a new data key...
    GenerateDataKeyRequest dataKeyRequest = new GenerateDataKeyRequest();
    dataKeyRequest.setKeyId(keyId);
    dataKeyRequest.setKeySpec("AES_128");

    GenerateDataKeyResult dataKeyResult = kms.generateDataKey(dataKeyRequest);

    ByteBuffer plaintextKey = dataKeyResult.getPlaintext();
    ByteBuffer encryptedKey = dataKeyResult.getCiphertextBlob();

    try {
        byte[] originalEncrypted = AES.encrypt(AES.pack(dataIn), plaintextKey.array());
        //Joining the encryptedKey and the encrypted bytes
        byte[] encryptedKeyBytes = encryptedKey.array();
        byte[] destination = new byte[originalEncrypted.length + encryptedKeyBytes.length];
        System.arraycopy(encryptedKeyBytes, 0, destination, 0, encryptedKeyBytes.length);
        System.arraycopy(originalEncrypted, 0, destination, encryptedKeyBytes.length, originalEncrypted.length);

        return AES.encode(destination);
    } catch (Exception e) {
        e.printStackTrace();
        Fault fault = new Fault();
        fault.setErrorCode(500);
        throw new FaultResponse("Error doing AES encryption: " + e.getLocalizedMessage(), fault);
    }
}

/**
 * @param dataIn
 * @return returns byte[]
 * @throws FaultResponse
 */
@Override
public String protectGenericData(String dataIn) throws FaultResponse {
    //Get a new data key...
    GenerateDataKeyRequest dataKeyRequest = new GenerateDataKeyRequest();
    dataKeyRequest.setKeyId(keyId);
    dataKeyRequest.setKeySpec("AES_128");

    GenerateDataKeyResult dataKeyResult = kms.generateDataKey(dataKeyRequest);

    ByteBuffer plaintextKey = dataKeyResult.getPlaintext();
    ByteBuffer encryptedKey = dataKeyResult.getCiphertextBlob();

    try {
        byte[] originalEncrypted = AES.encrypt(AES.pack(dataIn), plaintextKey.array());
        //Joining the encryptedKey and the encrypted bytes
        byte[] encryptedKeyBytes = encryptedKey.array();
        byte[] destination = new byte[originalEncrypted.length + encryptedKeyBytes.length];
        System.arraycopy(encryptedKeyBytes, 0, destination, 0, encryptedKeyBytes.length);
        System.arraycopy(originalEncrypted, 0, destination, encryptedKeyBytes.length, originalEncrypted.length);

        return AES.encode(destination);
    } catch (Exception e) {
        e.printStackTrace();
        Fault fault = new Fault();
        fault.setErrorCode(500);
        throw new FaultResponse("Error doing AES encryption: " + e.getLocalizedMessage(), fault);
    }
}

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

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

@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 = kms.generateDataKey(req);

    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.encodeAsString(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);
}