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

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

@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 Observable<Encrypted> encrypt(VertxContext<Server> vertxContext, byte[] plainBytes) {
    SfsVertx sfsVertx = vertxContext.vertx();
    Context context = sfsVertx.getOrCreateContext();
    return Observable.defer(() -> {
        byte[] cloned = Arrays.copyOf(plainBytes, plainBytes.length);
        return RxHelper.executeBlocking(context, 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);
            }
        });
    });
}

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

}

@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 testEncrypt_shouldThrowExceptionIfRequestFailed() throws Exception {
  AmazonServiceException exception = mock(AmazonServiceException.class);
  when(mockKms.encrypt(isA(EncryptRequest.class)))
      .thenThrow(exception);

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

@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.
  }
}

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

private EncryptResult encrypt0(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);

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

/**
 * Using the given AWS Key, encrypt the given bytes
 *
 * @param awsKeyId   unique identifier for the customer master key
 * @param clearBytes the unencrypted bytes to encrypt
 * @return the encrypted bytes
 */
public static byte[] encrypt(String awsKeyId, Map<String, String> encryptionContext, byte[] clearBytes) {
    EncryptRequest encryptRequest = new EncryptRequest();
    encryptRequest.setKeyId(awsKeyId);
    encryptRequest.setPlaintext(ByteBuffer.wrap(clearBytes));
    encryptRequest.setEncryptionContext(encryptionContext);

    AWSKMSClient client = new AWSKMSClient();
    EncryptResult encryptResult = client.encrypt(encryptRequest);

    return encryptResult.getCiphertextBlob().array();
}

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

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

        return extractString(ByteBuffer.wrap(Base64.encode(encryptedBytes.array())));
    }
}

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

@Override
public byte[] encrypt(final byte[] plaintext, final byte[] associatedData)
    throws GeneralSecurityException {
  try {
    EncryptRequest req =
        new EncryptRequest().withKeyId(keyArn).withPlaintext(ByteBuffer.wrap(plaintext));
    if (associatedData != null && associatedData.length != 0) {
      req = req.addEncryptionContextEntry("associatedData", BinaryUtils.toHex(associatedData));
    }
    return kmsClient.encrypt(req).getCiphertextBlob().array();
  } catch (AmazonServiceException e) {
    throw new GeneralSecurityException("encryption failed", e);
  }
}

public EncryptResult encrypt(EncryptRequest request) {
    // Default AWS limit was 1200 shared as of Aug 2017
    return execute("KmsEncryptDecrypt", "KmsEncrypt", () -> client.encrypt(request));
}

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

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