Merge pull request #22 from alex-litvak/master

Support PEM and plaintext formats for public key

src/main/java/io/iron/ironworker/client/Client.java

@@ -9,6 +9,7 @@
 import io.iron.ironworker.client.entities.*;
 import org.apache.http.HttpHost;
 
+import java.io.File;
 import java.io.IOException;
 import java.security.GeneralSecurityException;
 import java.util.ArrayList;
@@ -278,11 +279,14 @@ public void reload(ScheduleEntity entity) throws APIException {
     private String loadPayload(Map<String, Object> params, Map<String, Object> options) throws GeneralSecurityException, IOException {
         String payload = gson.toJson(params);
         if (options != null) {
-            String encryptionKeyFile = (String) options
-                    .get("encryptionKeyFile");
+            String encryptionKeyFile = (String) options.get("encryptionKeyFile");
+            String encryptionKey = (String) options.get("encryptionKey");
             if (encryptionKeyFile != null) {
-                PayloadEncryptor payloadEncryptor = new PayloadEncryptor(encryptionKeyFile);
-                payload = payloadEncryptor.encryptPayload(payload);
+                PayloadEncryptor payloadEncryptor = new PayloadEncryptor(new File(encryptionKeyFile));
+                payload = payloadEncryptor.encrypt(payload);
+            } else if(encryptionKey != null) {
+                PayloadEncryptor payloadEncryptor = new PayloadEncryptor(encryptionKey);
+                payload = payloadEncryptor.encrypt(payload);
             }
         }
         return payload;

src/main/java/io/iron/ironworker/client/builders/ScheduleOptions.java

@@ -35,6 +35,10 @@ public static ScheduleOptionsObject label(String label) {
         return (new ScheduleOptionsObject()).label(label);
     }
 
+    public ScheduleOptionsObject encryptionKey(String encryptionKey){
+        return (new ScheduleOptionsObject()).encryptionKey(encryptionKey);
+    }
+
     public ScheduleOptionsObject encryptionKeyFile(String encryptionKeyFile){
         return (new ScheduleOptionsObject()).encryptionKeyFile(encryptionKeyFile);
     }

src/main/java/io/iron/ironworker/client/builders/ScheduleOptionsObject.java

@@ -59,6 +59,12 @@ public ScheduleOptionsObject label(String label) {
         return this;
     }
 
+    public ScheduleOptionsObject encryptionKey(String encryptionKey) {
+        options.put("encryptionKey", encryptionKey);
+
+        return this;
+    }
+
     public ScheduleOptionsObject encryptionKeyFile(String encryptionKeyFile) {
         options.put("encryptionKeyFile", encryptionKeyFile);
 

src/main/java/io/iron/ironworker/client/builders/TaskOptions.java

@@ -21,6 +21,10 @@ public static TaskOptionsObject delay(int delay) {
         return (new TaskOptionsObject()).delay(delay);
     }
 
+    public static TaskOptionsObject encryptionKey(String encryptionKey){
+        return (new TaskOptionsObject()).encryptionKey(encryptionKey);
+    }
+
     public static TaskOptionsObject encryptionKeyFile(String encryptionKeyFile){
         return (new TaskOptionsObject()).encryptionKeyFile(encryptionKeyFile);
     }

src/main/java/io/iron/ironworker/client/builders/TaskOptionsObject.java

@@ -39,7 +39,13 @@ public TaskOptionsObject delay(int delay) {
 
         return this;
     }
-
+
+    public TaskOptionsObject encryptionKey(String encryptionKey) {
+        options.put("encryptionKey", encryptionKey);
+
+        return this;
+    }
+
     public TaskOptionsObject encryptionKeyFile(String encryptionKeyFile) {
         options.put("encryptionKeyFile", encryptionKeyFile);
 

src/main/java/io/iron/ironworker/client/encryptors/Algorithm.java

@@ -0,0 +1,8 @@
+package io.iron.ironworker.client.encryptors;
+
+public final class Algorithm {
+    public static final String RSA = "RSA";
+    public static final String AES = "AES";
+    public static final String AES_GCM = "AES/GCM/NoPadding";
+    public static final String OAEP_WITH_SHA1 = "RSA/ECB/OAEPWithSHA1AndMGF1Padding";
+}

src/main/java/io/iron/ironworker/client/encryptors/PayloadEncryptor.java

@@ -4,6 +4,7 @@
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.io.StringReader;
 import java.security.GeneralSecurityException;
 import java.security.Key;
 import java.security.KeyFactory;
@@ -15,28 +16,38 @@
 import javax.crypto.KeyGenerator;
 import javax.crypto.spec.GCMParameterSpec;
 
+import org.apache.commons.io.FileUtils;
+import org.apache.commons.io.FilenameUtils;
 import org.apache.commons.lang3.ArrayUtils;
 import org.bouncycastle.util.encoders.Base64;
+import org.bouncycastle.util.io.pem.PemObject;
+import org.bouncycastle.util.io.pem.PemReader;
 
 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 static final String PEM_FILE_EXTENSION = "pem";
+    private static final String DER_FILE_EXTENSION = "der";
 
     private PublicKey rsaPublicKey;
     private Key aesKey;
     private SecureRandom secureRandom;
 
-    public PayloadEncryptor(String encryptionKeyFile) throws GeneralSecurityException, IOException {
-        Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider());
-        secureRandom = new SecureRandom();
+    public PayloadEncryptor(String encryptionKey) throws GeneralSecurityException, IOException {
+        init();
+        rsaPublicKey = getPublicKey(encryptionKey);
+    }
+
+    public PayloadEncryptor(File encryptionKeyFile) throws GeneralSecurityException, IOException {
+        init();
         rsaPublicKey = getPublicKey(encryptionKeyFile);
-        aesKey = generateAESKey();
     }
 
-    public String encryptPayload(String payload) throws GeneralSecurityException {
+    public String encrypt(String payload) throws GeneralSecurityException {
         byte[] aesKeyBytes = aesKey.getEncoded();
         byte[] aesKeyCipher = encryptRSA(aesKeyBytes, rsaPublicKey);
 
@@ -49,18 +60,62 @@ public String encryptPayload(String payload) throws GeneralSecurityException {
         System.arraycopy(cipherPayload, 0, ciphertext, 0, cipherPayload.length);
         return new String(Base64.encode(ArrayUtils.addAll(aesKeyCipher, ciphertext)));
     }
+
+    private void init() throws GeneralSecurityException{
+        Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider());
+        secureRandom = new SecureRandom();
+        aesKey = generateAESKey();
+    }
+
+    private PublicKey getPublicKey(String encryptionKey) throws IOException, GeneralSecurityException{
+        PemReader pemReader = new PemReader(new StringReader(encryptionKey));
+        PemObject pemObject = pemReader.readPemObject();
+        pemReader.close();
+        byte[] keyBytes = pemObject.getContent();
+        return generatePublicKeyFromBytes(keyBytes);
+    }
+
+    private PublicKey getPublicKey(File encryptionKeyFile) throws GeneralSecurityException, IOException {
+        byte[] keyBytes;
+        PublicKey publicKey;
+        String keyFileExtension = FilenameUtils.getExtension(encryptionKeyFile.getName());
+        if(keyFileExtension.equals(PEM_FILE_EXTENSION)){
+            String publicKeyString = FileUtils.readFileToString(encryptionKeyFile);
+            publicKey = getPublicKey(publicKeyString);
+        } else if (keyFileExtension.equals(DER_FILE_EXTENSION)){
+            keyBytes = parseDerFile(encryptionKeyFile);
+            publicKey = generatePublicKeyFromBytes(keyBytes);
+        } else {
+            throw new UnsupportedKeyFileFormatException("Unsupported key file format: "+ keyFileExtension);
+        }
+
+        return publicKey;
+    }    
+
+    private byte[] parseDerFile(File derFile) throws IOException {
+        FileInputStream fis = new FileInputStream(derFile);
+        byte[] keyBytes = new byte[(int) derFile.length()];
+        DataInputStream dis = new DataInputStream(fis);
+        dis.readFully(keyBytes);
+        dis.close();
+        return keyBytes;
+    }
+
+    private PublicKey generatePublicKeyFromBytes(byte[] keyBytes) throws GeneralSecurityException {
+        X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
+        KeyFactory kf = KeyFactory.getInstance(Algorithm.RSA);
+        return kf.generatePublic(spec);
+    }
 
     private Key generateAESKey() throws GeneralSecurityException {
-        KeyGenerator generator;
-        generator = KeyGenerator.getInstance("AES", "SunJCE");
+        KeyGenerator generator = KeyGenerator.getInstance(Algorithm.AES);
         generator.init(AES_KEY_SIZE);
-        Key keyToBeWrapped = generator.generateKey();
-        return keyToBeWrapped;
+        Key aesKey = generator.generateKey();
+        return aesKey;
     }
 
     private byte[] encryptRSA(byte[] input, Key publicKey) throws GeneralSecurityException {
-        Cipher cipher;
-        cipher = Cipher.getInstance("RSA/ECB/OAEPWithSHA1AndMGF1Padding");
+        Cipher cipher = Cipher.getInstance(Algorithm.OAEP_WITH_SHA1);
         cipher.init(Cipher.ENCRYPT_MODE, publicKey, secureRandom);
         return encrypt(input, cipher);
     }
@@ -72,26 +127,11 @@ private byte[] encrypt(byte[] input, Cipher cipher) throws GeneralSecurityExcept
     }
 
     private Cipher getAesCipher(Key aesKey) throws GeneralSecurityException {
-        Cipher cipher;
         byte[] nonce = new byte[GCM_NONCE_LENGTH];
-        cipher = Cipher.getInstance("AES/GCM/NoPadding", "BC");
+        Cipher cipher = Cipher.getInstance(Algorithm.AES_GCM);        
         secureRandom.nextBytes(nonce);
         GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH * 8, nonce);
         cipher.init(Cipher.ENCRYPT_MODE, aesKey, spec);
-
         return cipher;
     }
-
-    private PublicKey getPublicKey(String filename) throws GeneralSecurityException, IOException {
-        File f = new File(filename);
-        FileInputStream fis;
-        fis = new FileInputStream(f);
-        DataInputStream dis = new DataInputStream(fis);
-        byte[] keyBytes = new byte[(int) f.length()];
-        dis.readFully(keyBytes);
-        dis.close();
-        X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);        
-        KeyFactory kf = KeyFactory.getInstance("RSA");
-        return kf.generatePublic(spec);        
-    }
 }

src/main/java/io/iron/ironworker/client/encryptors/UnsupportedKeyFileFormatException.java

@@ -0,0 +1,19 @@
+package io.iron.ironworker.client.encryptors;
+
+import org.bouncycastle.crypto.RuntimeCryptoException;
+
+public class UnsupportedKeyFileFormatException extends RuntimeCryptoException {
+    private String message;
+
+    public UnsupportedKeyFileFormatException(String message) {
+        this.message = message;
+    }
+
+    public String getMessage() {
+        return message;
+    }
+
+    public String toString() {
+        return message;
+    }
+}

src/test/java/io/iron/test/WorkerHelperTest.java

@@ -83,7 +83,6 @@ public void testParsingPayloadToString() throws IOException {
         String realPayload = helper.getPayload();
 
         System.out.println(realPayload);
-        String t = "{\"name\":\"test-name-3\",\"max_amount\":66,\"price\":99.99}";
         assertTrue(realPayload.matches(".*\"name\":\\s*\"test-name-3\".*"));
         assertTrue(realPayload.matches(".*\"max_amount\":\\s*66.*"));
         assertTrue(realPayload.matches(".*\"price\":\\s*99.99.*"));