Add opportunity to encrypt payload (RSA & AES)

iron-io · Feb 24, 2016 · 62a11c6 · 62a11c6
1 parent 8192d2c
commit 62a11c6
Show file tree

Hide file tree

Showing 6 changed files with 171 additions and 6 deletions.
diff --git a/pom.xml b/pom.xml
@@ -63,6 +63,16 @@
             <version>4.12</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+		    <groupId>org.bouncycastle</groupId>
+		    <artifactId>bcprov-jdk15on</artifactId>
+		    <version>1.54</version>
+  	    </dependency>
+  	    <dependency>
+	    	<groupId>commons-lang</groupId>
+			<artifactId>commons-lang</artifactId>
+			<version>2.3</version>
+    	</dependency>
     </dependencies>
 
     <profiles>

diff --git a/src/main/java/io/iron/ironworker/client/Client.java b/src/main/java/io/iron/ironworker/client/Client.java
@@ -5,6 +5,7 @@
 import com.google.gson.JsonObject;
 import io.iron.ironworker.client.builders.*;
 import io.iron.ironworker.client.codes.BaseCode;
+import io.iron.ironworker.client.encryptors.PayloadEncryptor;
 import io.iron.ironworker.client.entities.*;
 import org.apache.http.HttpHost;
 
@@ -123,9 +124,17 @@ public TaskEntity createTask(String codeName, Map<String, Object> params, Map<St
         if (params == null) {
             params = new HashMap<String, Object>();
         }
+        String payload = gson.toJson(params);
+        if(options!=null){
+        	String encryptionKeyFile = (String) options.get("encryptionKeyFile");
+            if (encryptionKeyFile!=null) {
+            	PayloadEncryptor payloadEncryptor = new PayloadEncryptor(encryptionKeyFile);
+            	payload = payloadEncryptor.encryptPayload(payload);
+            }
+        }
 
-        // TODO: implement multiple worker queueing (http://dev.iron.io/worker/reference/api/#queue_a_task)
-        TaskIds taskIds = gson.fromJson(api.tasksCreate(codeName, gson.toJson(params), options), TaskIds.class);
+        // TODO: implement multiple worker queueing (http://dev.iron.io/worker/reference/api/#queue_a_task)        
+        TaskIds taskIds = gson.fromJson(api.tasksCreate(codeName, payload, options), TaskIds.class);
         return taskIds.getTasks()[0];
     }
 
@@ -213,9 +222,17 @@ public ScheduleEntity createSchedule(String codeName, Map<String, Object> params
         if (params == null) {
             params = new HashMap<String, Object>();
         }
+        String payload = gson.toJson(params);
+        if(options!=null){
+        	String encryptionKeyFile = (String) options.get("encryptionKeyFile");
+            if (encryptionKeyFile!=null) {
+            	PayloadEncryptor payloadEncryptor = new PayloadEncryptor(encryptionKeyFile);
+            	payload = payloadEncryptor.encryptPayload(payload);
+            }
+        }
 
         // TODO: implement multiple worker scheduling (http://dev.iron.io/worker/reference/api/#schedule_a_task)
-        ScheduleIds scheduleIds = gson.fromJson(api.schedulesCreate(codeName, gson.toJson(params), options), ScheduleIds.class);
+        ScheduleIds scheduleIds = gson.fromJson(api.schedulesCreate(codeName, payload, options), ScheduleIds.class);
         return scheduleIds.getSchedules()[0];
     }
 

diff --git a/src/main/java/io/iron/ironworker/client/builders/TaskOptions.java b/src/main/java/io/iron/ironworker/client/builders/TaskOptions.java
@@ -20,6 +20,10 @@ public static TaskOptionsObject label(String label) {
     public static TaskOptionsObject delay(int delay) {
         return (new TaskOptionsObject()).delay(delay);
     }
+
+    public static TaskOptionsObject encryptionKeyFile(String encryptionKeyFile){
+    	return (new TaskOptionsObject()).encryptionKeyFile(encryptionKeyFile);
+    }
 
     protected TaskOptions() {
     }

diff --git a/src/main/java/io/iron/ironworker/client/builders/TaskOptionsObject.java b/src/main/java/io/iron/ironworker/client/builders/TaskOptionsObject.java
@@ -39,6 +39,12 @@ public TaskOptionsObject delay(int delay) {
 
         return this;
     }
+
+    public TaskOptionsObject encryptionKeyFile(String encryptionKeyFile){
+    	options.put("encryptionKeyFile", encryptionKeyFile);
+
+        return this;
+    }
 
     public Map<String, Object> create() {
         return options;

diff --git a/src/main/java/io/iron/ironworker/client/encryptors/PayloadEncryptor.java b/src/main/java/io/iron/ironworker/client/encryptors/PayloadEncryptor.java
@@ -0,0 +1,131 @@
+package io.iron.ironworker.client.encryptors;
+
+import io.iron.ironworker.client.APIException;
+import java.io.DataInputStream;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.security.Key;
+import java.security.KeyFactory;
+import java.security.PublicKey;
+import java.security.SecureRandom;
+import java.security.Security;
+import java.security.spec.X509EncodedKeySpec;
+import javax.crypto.Cipher;
+import javax.crypto.KeyGenerator;
+import javax.crypto.spec.GCMParameterSpec;
+
+import org.apache.commons.lang.*;
+import org.bouncycastle.util.encoders.Base64;
+
+public class PayloadEncryptor {
+
+	// AES-GCM parameters
+    private static final int AES_KEY_SIZE = 128; // in bits
+    private static final int GCM_NONCE_LENGTH = 12; // in bytes
+    private static final int GCM_TAG_LENGTH = 16; // in bytes (overhead)
+
+    private PublicKey rsaPublicKey;
+    private Key aesKey;
+    private SecureRandom secureRandom;
+
+    public PayloadEncryptor(String encryptionKeyFile) throws APIException {
+    	Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider());
+    	secureRandom = new SecureRandom();
+    	rsaPublicKey = getPublicKey(encryptionKeyFile);
+    	aesKey = generateAESKey();
+    }
+
+    public String encryptPayload(String payload) throws APIException {
+    	byte[] aesKeyBytes = aesKey.getEncoded();
+		byte[] aesKeyCipher = encryptRSA(aesKeyBytes, rsaPublicKey);
+
+        Cipher gcm = getAesCipher(aesKey);
+		byte[] pbytes = payload.getBytes();
+		byte[] cipherPayload = encrypt(pbytes, gcm);
+
+		byte[] ciphertext = new byte[pbytes.length + GCM_TAG_LENGTH + GCM_NONCE_LENGTH];
+		System.arraycopy(gcm.getIV(),0,ciphertext,ciphertext.length - GCM_NONCE_LENGTH, GCM_NONCE_LENGTH);
+		System.arraycopy(cipherPayload,0,ciphertext,0,cipherPayload.length);
+		return new String(Base64.encode(ArrayUtils.addAll(aesKeyCipher, ciphertext)));
+    }
+
+    private Key generateAESKey() throws APIException {
+	    KeyGenerator generator;
+		try {
+			generator = KeyGenerator.getInstance("AES", "SunJCE");
+		} catch (Exception e) {
+			throw new APIException(e.getMessage(), e);
+		}
+	    generator.init(AES_KEY_SIZE);
+	    Key keyToBeWrapped = generator.generateKey();
+	    return keyToBeWrapped;
+	}
+
+    private byte[] encryptRSA(byte[] input, Key publicKey) throws APIException {
+    	Cipher cipher;
+	    try {
+	    	cipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA1AndMGF1Padding");
+			cipher.init(Cipher.ENCRYPT_MODE, publicKey, secureRandom);
+		} catch (Exception e) {
+			throw new APIException(e.getMessage(), e);
+		}	    
+	    return encrypt(input, cipher);
+	}
+
+    private byte[] encrypt(byte[] input, Cipher cipher) throws APIException {
+    	byte[] cipherText;
+		try {
+			cipherText = cipher.doFinal(input);
+		} catch (Exception e) {
+			throw new APIException(e.getMessage(), e);
+		}
+    	return cipherText;
+    }
+
+    private Cipher getAesCipher(Key aesKey) throws APIException {
+   		Cipher cipher;
+   		byte[] nonce = new byte[GCM_NONCE_LENGTH];
+		try {
+			cipher = Cipher.getInstance("AES/GCM/NoPadding", "BC");
+		} catch (Exception e) {
+			throw new APIException(e.getMessage(), e);
+		}
+        secureRandom.nextBytes(nonce);
+        GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH * 8, nonce);
+        try {
+			cipher.init(Cipher.ENCRYPT_MODE, aesKey, spec);
+		} catch (Exception e) {
+			throw new APIException(e.getMessage(), e);
+		}
+        return cipher;
+   	}
+
+    private PublicKey getPublicKey(String filename) throws APIException {
+	    File f = new File(filename);
+	    FileInputStream fis;
+		try {
+			fis = new FileInputStream(f);
+		} catch (FileNotFoundException e) {
+			throw new APIException(String.format("File %s not found", filename), e);
+		}
+	    DataInputStream dis = new DataInputStream(fis);
+	    byte[] keyBytes = new byte[(int)f.length()];
+	    try {
+			dis.readFully(keyBytes);
+		    dis.close();
+	    } catch (IOException e) {
+			throw new APIException(e.getMessage(), e);
+		}
+
+	    X509EncodedKeySpec spec =
+	      new X509EncodedKeySpec(keyBytes);	    
+		try {
+			KeyFactory kf = KeyFactory.getInstance("RSA");
+			return kf.generatePublic(spec);
+		} catch (Exception e) {
+			throw new APIException(e.getMessage(), e);
+		}
+    }
+}
diff --git a/src/main/java/io/iron/ironworker/client/helpers/WorkerHelper.java b/src/main/java/io/iron/ironworker/client/helpers/WorkerHelper.java
@@ -2,8 +2,6 @@
 
 import com.google.gson.Gson;
 import com.google.gson.JsonObject;
-import org.apache.commons.io.IOUtils;
-
 import java.io.*;
 import java.nio.file.Files;
 import java.nio.file.Paths;
@@ -88,7 +86,6 @@ public String getPayload() throws IOException {
             return payload;
         }
 
-        FileInputStream inputStream = new FileInputStream(payloadPath);
         payload = new String(Files.readAllBytes(Paths.get(payloadPath)));
         return payload;
     }