KMS simple test

src/main/java/cloud/localstack/Localstack.java

@@ -139,6 +139,10 @@ public int getEdgePort() {
     public String getEndpointKinesis() {
         return endpointForService(ServiceName.KINESIS);
     }
+
+    public String getEndpointKMS() {
+        return endpointForService(ServiceName.KMS);
+    }
 
     public String getEndpointLambda() {
         return endpointForService(ServiceName.LAMBDA);

src/main/java/cloud/localstack/ServiceName.java

@@ -3,6 +3,7 @@
 public class ServiceName {
     public static final String API_GATEWAY = "apigateway";
     public static final String KINESIS = "kinesis";
+    public static final String KMS = "kms";
     public static final String DYNAMO = "dynamodb";
     public static final String DYNAMO_STREAMS = "dynamodbstreams";
     public static final String ELASTICSEARCH = "elasticsearch";

src/main/java/cloud/localstack/awssdkv1/PowerMockLocalStack.java

@@ -22,6 +22,10 @@
 import com.amazonaws.services.kinesis.AmazonKinesisAsync;
 import com.amazonaws.services.kinesis.AmazonKinesisAsyncClientBuilder;
 import com.amazonaws.services.kinesis.AmazonKinesisClientBuilder;
+import com.amazonaws.services.kms.AWSKMS;
+import com.amazonaws.services.kms.AWSKMSAsync;
+import com.amazonaws.services.kms.AWSKMSAsyncClientBuilder;
+import com.amazonaws.services.kms.AWSKMSClientBuilder;
 import com.amazonaws.services.lambda.AWSLambda;
 import com.amazonaws.services.lambda.AWSLambdaAsync;
 import com.amazonaws.services.lambda.AWSLambdaAsyncClientBuilder;
@@ -57,6 +61,7 @@
         AWSLambdaClientBuilder.class, AWSLogsClientBuilder.class, AWSLogsAsyncClientBuilder.class,
         AmazonS3ClientBuilder.class, AmazonSimpleEmailServiceAsyncClientBuilder.class,
         AmazonSimpleEmailServiceClientBuilder.class, AmazonSNSAsyncClientBuilder.class, AmazonSNSClientBuilder.class,
+        AWSKMSClientBuilder.class, AWSKMSAsyncClientBuilder.class,
         AWSSecretsManagerClientBuilder.class, AmazonSQSClientBuilder.class })
 @PowerMockIgnore({ "javax.crypto.*", "org.hamcrest.*", "javax.net.ssl.*", "com.sun.org.apache.xerces.*", "javax.xml.*",
         "org.xml.*", "javax.management.*", "javax.security.*", "org.w3c.*" })
@@ -97,6 +102,18 @@ public static void mockKinesis() {
         PowerMockito.mockStatic(AmazonKinesisClientBuilder.class);
         when(AmazonKinesisClientBuilder.defaultClient()).thenReturn(mockService);
     }
+
+    public static void mockKMS() {
+        AWSKMS mockService = TestUtils.getClientKMS();
+        PowerMockito.mockStatic(AWSKMSClientBuilder.class);
+        when(AWSKMSClientBuilder.defaultClient()).thenReturn(mockService);
+    }
+
+    public static void mockKMSAsync() {
+        AWSKMSAsync mockService = TestUtils.getClientKMSAsync();
+        PowerMockito.mockStatic(AWSKMSAsyncClientBuilder.class);
+        when(AWSKMSAsyncClientBuilder.defaultClient()).thenReturn(mockService);
+    }
 
     public static void mockKinesisAsync() {
         AmazonKinesisAsync mockServiceAsync = TestUtils.getClientKinesisAsync();

src/main/java/cloud/localstack/awssdkv1/TestUtils.java

@@ -20,6 +20,10 @@
 import com.amazonaws.services.kinesis.AmazonKinesisAsync;
 import com.amazonaws.services.kinesis.AmazonKinesisAsyncClientBuilder;
 import com.amazonaws.services.kinesis.AmazonKinesisClientBuilder;
+import com.amazonaws.services.kms.AWSKMS;
+import com.amazonaws.services.kms.AWSKMSAsync;
+import com.amazonaws.services.kms.AWSKMSAsyncClientBuilder;
+import com.amazonaws.services.kms.AWSKMSClientBuilder;
 import com.amazonaws.services.lambda.AWSLambda;
 import com.amazonaws.services.lambda.AWSLambdaAsync;
 import com.amazonaws.services.lambda.AWSLambdaAsyncClientBuilder;
@@ -203,6 +207,17 @@ public static AmazonIdentityManagementAsync getClientIAMAsync() {
                 withCredentials(getCredentialsProvider()).build();
     }
 
+    public static AWSKMS getClientKMS(){
+        return AWSKMSClientBuilder.standard().
+            withEndpointConfiguration(getEndpointConfigurationKMS()).
+            withCredentials(getCredentialsProvider()).build();
+    }
+
+    public static AWSKMSAsync getClientKMSAsync(){
+        return AWSKMSAsyncClientBuilder.standard().
+            withEndpointConfiguration(getEndpointConfigurationKMS()).
+            withCredentials(getCredentialsProvider()).build();
+    }
 
     protected static AwsClientBuilder.EndpointConfiguration getEndpointConfigurationIAM() {
         return getEndpointConfiguration(Localstack.INSTANCE.getEndpointIAM());
@@ -255,6 +270,10 @@ protected static AwsClientBuilder.EndpointConfiguration getEndpointConfiguration
     protected static AwsClientBuilder.EndpointConfiguration getEndpointConfigurationStepFunctions() {
         return getEndpointConfiguration(Localstack.INSTANCE.getEndpointStepFunctions());
     }
+
+    protected static AwsClientBuilder.EndpointConfiguration getEndpointConfigurationKMS() {
+        return getEndpointConfiguration(Localstack.INSTANCE.getEndpointKMS());
+    }
 
     /**
      * UTIL METHODS

src/test/java/cloud/localstack/KMSTest.java

@@ -0,0 +1,70 @@
+package cloud.localstack;
+
+import cloud.localstack.awssdkv1.PowerMockLocalStack;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+import java.nio.ByteBuffer;
+
+import com.amazonaws.services.kms.AWSKMS;
+import com.amazonaws.services.kms.AWSKMSClientBuilder;
+import com.amazonaws.services.kms.model.CreateKeyRequest;
+import com.amazonaws.services.kms.model.CreateKeyResult;
+import com.amazonaws.services.kms.model.DecryptRequest;
+import com.amazonaws.services.kms.model.DecryptResult;
+import com.amazonaws.services.kms.model.EncryptRequest;
+import com.amazonaws.services.kms.model.EncryptResult;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+/**
+ * Test integration of KMS with LocalStack
+ */
+public class KMSTest extends PowerMockLocalStack {
+    private AWSKMS awskms;
+    private String keyId = "";
+
+    @Before
+    public void mock() {
+        PowerMockLocalStack.mockKMS();
+        awskms = AWSKMSClientBuilder.defaultClient();
+    }
+
+    @Test
+    public void createKeyTest() throws Exception {
+        CreateKeyRequest request = new CreateKeyRequest().withDescription("test key");
+        CreateKeyResult result = awskms.createKey(request);
+        keyId = result.getKeyMetadata().getKeyId();
+        Assert.assertNotNull(keyId);
+    }
+
+    @Test
+    public void stringEncriptionTest() throws Exception {
+        if(keyId.isEmpty()){
+            createKeyTest();
+        }
+
+        String testMessage =  "hey, this is a test of encryption";
+        byte[] encodedMessage = encrypt(testMessage.getBytes());
+        byte[] decodedMessage = decrypt(encodedMessage);
+        String resultMessage = new String(decodedMessage);
+
+        Assert.assertEquals(testMessage, resultMessage);
+    }
+
+    public byte[] encrypt(byte[] input) {
+        ByteBuffer buffer = ByteBuffer.wrap(input);
+        EncryptResult result = awskms.encrypt(new EncryptRequest().withKeyId(keyId).withPlaintext(buffer));
+        return result.getCiphertextBlob().array();
+    }
+
+    public byte[] decrypt(byte[] input) {
+        ByteBuffer buffer = ByteBuffer.wrap(input);
+        DecryptResult result = awskms.decrypt(new DecryptRequest().withCiphertextBlob(buffer));
+        return result.getPlaintext().array();
+    }
+
+}