Java 类com.amazonaws.encryptionsdk.kms.KmsMasterKeyProvider 实例源码

@BeforeMethod
public void setUp() throws Exception {
    AWSCredentialsProvider mockCredentials = mock(AWSCredentialsProvider.class);
    ClientConfiguration mockConfig = mock(ClientConfiguration.class);
    SecretsGroupIdentifier group = new SecretsGroupIdentifier(Region.US_WEST_1, "test.group");

    this.mockAwsCrypto = mock(AwsCrypto.class);
    this.mockKmsManager = mock(KMSManager.class);
    KMSEncryptor encryptor = new KMSEncryptor(mockKmsManager, mockCredentials, mockConfig, group, mockAwsCrypto, EncryptionStrength.AES_256);

    this.kmsEncryptor = spy(encryptor);
    this.mockProvider = mock(KmsMasterKeyProvider.class);
    doReturn(mockProvider).when(kmsEncryptor).getProvider();

    // Verify the expected encryption algorithm was set.
    verify(mockAwsCrypto, times(1)).setEncryptionAlgorithm(
            CryptoAlgorithm.ALG_AES_256_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384);
}

@Test
public void whenConstructedWithoutArguments_canUseMultipleRegions() throws Exception {
    KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder().build();

    for (String key : KMSTestFixtures.TEST_KEY_IDS) {
        byte[] ciphertext =
                new AwsCrypto().encryptData(
                        KmsMasterKeyProvider.builder()
                            .withKeysForEncryption(key)
                            .build(),
                        new byte[1]
                ).getResult();

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

@SuppressWarnings("deprecation") @Test(expected = CannotUnwrapDataKeyException.class)
public void whenLegacyConstructorsUsed_multiRegionDecryptIsNotSupported() throws Exception {
    KmsMasterKeyProvider mkp = new KmsMasterKeyProvider();

    for (String key : KMSTestFixtures.TEST_KEY_IDS) {
        byte[] ciphertext =
                new AwsCrypto().encryptData(
                        KmsMasterKeyProvider.builder()
                                            .withKeysForEncryption(key)
                                            .build(),
                        new byte[1]
                ).getResult();

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

@Test
public void whenHandlerConfigured_handlerIsInvoked() throws Exception {
    RequestHandler2 handler = spy(new RequestHandler2() {});
    KmsMasterKeyProvider mkp =
            KmsMasterKeyProvider.builder()
                                .withClientBuilder(
                                        AWSKMSClientBuilder.standard()
                                            .withRequestHandlers(handler)
                                )
                                .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
                                .build();

    new AwsCrypto().encryptData(mkp, new byte[1]);

    verify(handler).beforeRequest(any());
}

@Test
public void whenCustomCredentialsSet_theyAreUsed() throws Exception {
    AWSCredentialsProvider customProvider = spy(new DefaultAWSCredentialsProviderChain());

    KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder()
                                                   .withCredentials(customProvider)
                                                   .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
                                                   .build();

    new AwsCrypto().encryptData(mkp, new byte[1]);

    verify(customProvider, atLeastOnce()).getCredentials();

    AWSCredentials customCredentials = spy(customProvider.getCredentials());

    mkp = KmsMasterKeyProvider.builder()
                                                   .withCredentials(customCredentials)
                                                   .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
                                                   .build();

    new AwsCrypto().encryptData(mkp, new byte[1]);

    verify(customCredentials, atLeastOnce()).getAWSSecretKey();
}

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

/**
 * {@inheritDoc }
 */
@Override
public void decryptFile(
        final String encryptedFilename, 
        final String decryptedFilename) {

    final KmsMasterKeyProvider provider
            = new KmsMasterKeyProvider(
                    new DefaultAWSCredentialsProviderChain());

    final AwsCrypto awsCrypto
            = new AwsCrypto();

    try (final FileInputStream fileInputStream
            = new FileInputStream(
                    encryptedFilename);

            final FileOutputStream fileOutputStream
                    = new FileOutputStream(
                            decryptedFilename);

            final CryptoInputStream<?> decryptingStream
                    = awsCrypto
                            .createDecryptingStream(
                                    provider, 
                                    fileInputStream)) {

        IOUtils.copy(
                decryptingStream,
                fileOutputStream);

    } catch (IOException exception) {
        throw new DecryptionException(exception);
    }
}

/**
 * {@inheritDoc }
 */
@Override
public String decryptFile(
        final String encryptedFilename) {

    final KmsMasterKeyProvider provider
            = new KmsMasterKeyProvider(
                    new DefaultAWSCredentialsProviderChain());

    final AwsCrypto awsCrypto
            = new AwsCrypto();

    try (final FileInputStream fileInputStream
            = new FileInputStream(
                    encryptedFilename);

            final CryptoInputStream<?> decryptingStream
                    = awsCrypto
                            .createDecryptingStream(
                                    provider, 
                                    fileInputStream)) {

        return IOUtils.toString(
                decryptingStream);

    } catch (IOException exception) {
        throw new DecryptionException(exception);
    }
}

private MasterKeyProvider<?> masterKeyProvider() {

        final AWSCredentialsProvider credentials
                = new DefaultAWSCredentialsProviderChain();

        List<KmsMasterKey> masterKeys
                = new LinkedList<>();

        for (String region : this.regions) {
            KmsMasterKeyProvider provider
                    = new KmsMasterKeyProvider(
                            credentials,
                            Region.getRegion(
                                    Regions.fromName(
                                            region)),
                            new ClientConfiguration(),
                            this.keyId);



            masterKeys.add(
                    provider.getMasterKey(
                            this.keyId));
        }

        return MultipleProviderFactory
                .buildMultiProvider(
                        masterKeys);
    }

protected KmsMasterKeyProvider getProvider() {
    if (!prov.isPresent()) {
        Region region = RegionUtils.getRegion(groupIdentifier.region.getName());
        prov = Optional.of(new KmsMasterKeyProvider(awsCredentials, region, transformAndVerifyOrThrow(clientConfiguration), getKeyArn()));
    }
    return prov.get();
}

private static void standardEncrypt(final String kmsArn, final String fileName) throws Exception {
    // Encrypt with the KMS CMK and the escrowed public key
    // 1. Instantiate the SDK
    final AwsCrypto crypto = new AwsCrypto();

    // 2. Instantiate a KMS master key provider
    final KmsMasterKeyProvider kms = new KmsMasterKeyProvider(kmsArn);

    // 3. Instantiate a JCE master key provider
    // Because the user does not have access to the private escrow key,
    // they pass in "null" for the private key parameter.
    final JceMasterKey escrowPub = JceMasterKey.getInstance(publicEscrowKey, null, "Escrow", "Escrow",
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding");

    // 4. Combine the providers into a single master key provider
    final MasterKeyProvider<?> provider = MultipleProviderFactory.buildMultiProvider(kms, escrowPub);

    // 5. Encrypt the file
    // To simplify the code, we omit the encryption context. Production code should always 
    // use an encryption context. For an example, see the other SDK samples.
    final FileInputStream in = new FileInputStream(fileName);
    final FileOutputStream out = new FileOutputStream(fileName + ".encrypted");
    final CryptoOutputStream<?> encryptingStream = crypto.createEncryptingStream(provider, out);

    IOUtils.copy(in, encryptingStream);
    in.close();
    encryptingStream.close();
}

private static void standardDecrypt(final String kmsArn, final String fileName) throws Exception {
    // Decrypt with the KMS CMK and the escrow public key. You can use a combined provider, 
    // as shown here, or just the KMS master key provider.

    // 1. Instantiate the SDK
    final AwsCrypto crypto = new AwsCrypto();

    // 2. Instantiate a KMS master key provider
    final KmsMasterKeyProvider kms = new KmsMasterKeyProvider(kmsArn);

    // 3. Instantiate a JCE master key provider
    // Because the user does not have access to the private 
    // escrow key, they pass in "null" for the private key parameter.
    final JceMasterKey escrowPub = JceMasterKey.getInstance(publicEscrowKey, null, "Escrow", "Escrow",
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding");

    // 4. Combine the providers into a single master key provider
    final MasterKeyProvider<?> provider = MultipleProviderFactory.buildMultiProvider(kms, escrowPub);

    // 5. Decrypt the file
    // To simplify the code, we omit the encryption context. Production code should always 
    // use an encryption context. For an example, see the other SDK samples.
    final FileInputStream in = new FileInputStream(fileName + ".encrypted");
    final FileOutputStream out = new FileOutputStream(fileName + ".decrypted");
    final CryptoOutputStream<?> decryptingStream = crypto.createDecryptingStream(provider, out);
    IOUtils.copy(in, decryptingStream);
    in.close();
    decryptingStream.close();
}

@Test
public void testMultipleRegionKmsKeys() {
    final MockKMSClient us_east_1 = new MockKMSClient();
    us_east_1.setRegion(Region.getRegion(Regions.US_EAST_1));
    final MockKMSClient eu_west_1 = new MockKMSClient();
    eu_west_1.setRegion(Region.getRegion(Regions.EU_WEST_1));
    final String arn1 = us_east_1.createKey().getKeyMetadata().getArn();
    final String arn2 = eu_west_1.createKey().getKeyMetadata().getArn();
    KmsMasterKeyProvider provE = legacyConstruct(us_east_1, Region.getRegion(Regions.US_EAST_1));
    KmsMasterKeyProvider provW = legacyConstruct(eu_west_1, Region.getRegion(Regions.EU_WEST_1));
    KmsMasterKey mk1 = provE.getMasterKey(arn1);
    KmsMasterKey mk2 = provW.getMasterKey(arn2);

    final MasterKeyProvider<KmsMasterKey> mkp = MultipleProviderFactory.buildMultiProvider(KmsMasterKey.class,
                                                                                           mk1, mk2);
    AwsCrypto crypto = new AwsCrypto();
    CryptoResult<byte[], KmsMasterKey> ct = crypto.encryptData(mkp, PLAINTEXT);
    assertEquals(2, ct.getMasterKeyIds().size());

    CryptoResult<byte[], KmsMasterKey> result = crypto.decryptData(mk1, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk1, result.getMasterKeys().get(0));

    result = crypto.decryptData(mk2, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));

    assertMultiReturnsKeys(mkp, mk1, mk2);

    // Delete one of the two keys and ensure it's still decryptable
    us_east_1.deleteKey(arn1);

    result = crypto.decryptData(mkp, ct.getResult());
    assertArrayEquals(PLAINTEXT, result.getResult());
    // Only the first found key should be used
    assertEquals(1, result.getMasterKeys().size());
    assertEquals(mk2, result.getMasterKeys().get(0));
}

@Test
public void whenShortTimeoutSet_timesOut() throws Exception {
    // By setting a timeout of 1ms, it's not physically possible to complete both the us-west-2 and eu-central-1
    // requests due to speed of light limits.
    KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder()
                                                   .withClientBuilder(
                                                           AWSKMSClientBuilder.standard()
                                                            .withClientConfiguration(
                                                                    new ClientConfiguration()
                                                                        .withRequestTimeout(1)
                                                            )
                                                   )
                                                   .withKeysForEncryption(Arrays.asList(KMSTestFixtures.TEST_KEY_IDS))
                                                   .build();

    try {
        new AwsCrypto().encryptData(mkp, new byte[1]);
        fail("Expected exception");
    } catch (Exception e) {
        if (e instanceof AbortedException) {
            // ok - one manifestation of a timeout
        } else if (e.getCause() instanceof HttpRequestTimeoutException) {
            // ok - another kind of timeout
        } else {
            throw e;
        }
    }
}

@Test
public void whenBuilderCloned_credentialsAndConfigurationAreRetained() throws Exception {
    AWSCredentialsProvider customProvider1 = spy(new DefaultAWSCredentialsProviderChain());
    AWSCredentialsProvider customProvider2 = spy(new DefaultAWSCredentialsProviderChain());

    KmsMasterKeyProvider.Builder builder = KmsMasterKeyProvider.builder()
            .withCredentials(customProvider1)
            .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0]);

    KmsMasterKeyProvider.Builder builder2 = builder.clone();

    // This will mutate the first builder to add the new key and change the creds, but leave the clone unchanged.
    MasterKeyProvider<?> mkp2 = builder.withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[1]).withCredentials(customProvider2).build();
    MasterKeyProvider<?> mkp1 = builder2.build();

    CryptoResult<byte[], ?> result = new AwsCrypto().encryptData(mkp1, new byte[0]);

    assertEquals(KMSTestFixtures.TEST_KEY_IDS[0], result.getMasterKeyIds().get(0));
    assertEquals(1, result.getMasterKeyIds().size());
    verify(customProvider1, atLeastOnce()).getCredentials();
    verify(customProvider2, never()).getCredentials();

    reset(customProvider1, customProvider2);

    result = new AwsCrypto().encryptData(mkp2, new byte[0]);

    assertTrue(result.getMasterKeyIds().contains(KMSTestFixtures.TEST_KEY_IDS[0]));
    assertTrue(result.getMasterKeyIds().contains(KMSTestFixtures.TEST_KEY_IDS[1]));
    assertEquals(2, result.getMasterKeyIds().size());
    verify(customProvider1, never()).getCredentials();
    verify(customProvider2, atLeastOnce()).getCredentials();
}

@Test
public void whenBuilderCloned_clientBuilderCustomizationIsRetained() throws Exception {
    RequestHandler2 handler = spy(new RequestHandler2() {});

    KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder()
            .withClientBuilder(
                    AWSKMSClientBuilder.standard().withRequestHandlers(handler)
            )
            .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
            .clone().build();

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

    verify(handler, atLeastOnce()).beforeRequest(any());
}

@Test(expected = IllegalArgumentException.class)
public void whenBogusEndpointIsSet_constructionFails() throws Exception {
    KmsMasterKeyProvider.builder()
                        .withClientBuilder(
                                AWSKMSClientBuilder.standard()
                                                   .withEndpointConfiguration(
                                                           new AwsClientBuilder.EndpointConfiguration(
                                                                   "https://this.does.not.exist.example.com",
                                                                   "bad-region")
                                                   )
                        );
}

@Test
public void testDecryptFromFile() throws Exception {
    AwsCrypto crypto = new AwsCrypto();
    final KmsMasterKeyProvider masterKeyProvider = new KmsMasterKeyProvider(kmsKeyId);
    byte ciphertextBytes[] = Files.readAllBytes(Paths.get(ciphertextFileName));
    byte plaintextBytes[] = Files.readAllBytes(Paths.get(plaintextFileName));
    final CryptoResult decryptResult = crypto.decryptData(
        masterKeyProvider,
        ciphertextBytes
    );
    assertArrayEquals(plaintextBytes, (byte[])decryptResult.getResult());
}

public static void main(final String[] args) {
    keyArn = args[0];
    data = args[1];

    // Instantiate the SDK
    final AwsCrypto crypto = new AwsCrypto();

    // Set up the KmsMasterKeyProvider backed by the default credentials
    final KmsMasterKeyProvider prov = new KmsMasterKeyProvider(keyArn);

    // Encrypt the data
    //
    // Most encrypted data should have an associated encryption context
    // to protect integrity. This sample uses placeholder values.
    //
    // For more information see:
    // blogs.aws.amazon.com/security/post/Tx2LZ6WBJJANTNW/How-to-Protect-the-Integrity-of-Your-Encrypted-Data-by-Using-AWS-Key-Management
    final Map<String, String> context = Collections.singletonMap("Example", "String");

    final String ciphertext = crypto.encryptString(prov, data, context).getResult();
    System.out.println("Ciphertext: " + ciphertext);

    // Decrypt the data
    final CryptoResult<String, KmsMasterKey> decryptResult = crypto.decryptString(prov, ciphertext);

    // Before returning the plaintext, verify that the customer master key that
    // was used in the encryption operation was the one supplied to the master key provider. 
    if (!decryptResult.getMasterKeyIds().get(0).equals(keyArn)) {
        throw new IllegalStateException("Wrong key id!");
    }

    // Also, verify that the encryption context in the result contains the
    // encryption context supplied to the encryptString method. Because the
    // SDK can add values to the encryption context, don't require that 
    // the entire context matches. 
    for (final Map.Entry<String, String> e : context.entrySet()) {
        if (!e.getValue().equals(decryptResult.getEncryptionContext().get(e.getKey()))) {
            throw new IllegalStateException("Wrong Encryption Context!");
        }
    }

    // Now we can return the plaintext data
    System.out.println("Decrypted: " + decryptResult.getResult());
}

private KmsMasterKeyProvider legacyConstruct(final AWSKMS client, String... keyIds) {
    return legacyConstruct(client, Region.getRegion(Regions.DEFAULT_REGION), keyIds);
}

private KmsMasterKeyProvider legacyConstruct(final AWSKMS client, final Region region, String... keyIds) {
    return new KmsMasterKeyProvider(client, region, Arrays.asList(keyIds));
}

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