Code review

我們回頭再 review 一下週末專案，這兩天，讓我們一步一步把它擴充。

此外，我的程式，不太寫註解。程式的語言特性，不如 shell script，shell script 若不寫註解，大概只有寫的本人在寫的當下知道在寫什麼。但程式不是，現代的高階程式語言，都具有高度可自我文件化(self-documented)的特性。因此，可以利用名稱與函式的巧妙安排，讓程式不用註解也能順利閱讀。

專案簡介

一個簡單的 AES256 加密應用，用來加密檔案。

Test First

package javaxx.cipher;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.net.URISyntaxException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Arrays;

import org.apache.commons.io.IOUtils;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import static javaxx.cipher.EncryptionParameter.ALGORITHM_AES;
import static javaxx.cipher.EncryptionParameter.STRENGTH_256;

/**
 * AdvancedAesTest
 * cchuang, 2016/12/4.
 */
public class AdvancedAesTest {

    private static final String SOURCE_FILE = "/test.png";
    private static final String ENCRYPTED_FILE = "test.png.enc";
    private static final String DECRYPTED_FILE = "test.png.dec";
    private static final String TEST_PASSWORD = "s3cR3T p4SSw@Rd";

    private EncryptionParameter param;
    private File sourceFile;
    private File encryptedFile;
    private File decryptedFile;

    @Before
    public void setUp() throws Exception {
        param = initialEncryptionParameter();
        sourceFile = initSourceFile();
        encryptedFile = initEmptyFileBySource(sourceFile, ENCRYPTED_FILE);
        decryptedFile = initEmptyFileBySource(sourceFile, DECRYPTED_FILE);
    }

    private EncryptionParameter initialEncryptionParameter() throws NoSuchAlgorithmException {
        return EncryptionParameter.Builder.create()
                                          .setAlgorithm(ALGORITHM_AES)
                                          .setStrength(STRENGTH_256)
                                          .buildForEncryption();
    }

    private File initSourceFile() throws URISyntaxException {
        File fileInClasspath = new File(this.getClass().getResource(SOURCE_FILE).toURI());

        return fileInClasspath;
    }

    private File initEmptyFileBySource(File exists, String fileName) {
        File encryptedFile = exists.toPath()
                                   .getParent()
                                   .resolve(fileName)
                                   .toFile();

        return encryptedFile;
    }

    @Test
    public void testEncryptFile() throws Exception {
        EncryptionUtils.encryptFile(param, sourceFile, encryptedFile);

        checkEncryptedFile();

        EncryptionUtils.decryptFile(param, encryptedFile, decryptedFile);

        checkDecryptedFile();

        encryptedFile.delete();
    }

    private void checkEncryptedFile() throws IOException {
        Assert.assertTrue(encryptedFile.exists());
        Assert.assertFalse(
                Arrays.equals(IOUtils.toByteArray(new FileInputStream(sourceFile)),
                        IOUtils.toByteArray(new FileInputStream(encryptedFile))));
    }

    private void checkDecryptedFile() throws IOException {
        Assert.assertTrue(decryptedFile.exists());
        Assert.assertTrue(
                Arrays.equals(IOUtils.toByteArray(new FileInputStream(sourceFile)),
                        IOUtils.toByteArray(new FileInputStream(decryptedFile))));
    }

    @Test
    public void testEncryptSecretKeyByMasterKey() throws Exception {
        String masterKey = TEST_PASSWORD;
        byte[] salt = genSalt();

        EncryptionParameter encParam =
                EncryptionParameter.Builder.create()
                                           .setAlgorithm(ALGORITHM_AES)
                                           .setStrength(STRENGTH_256)
                                           .setPassword(masterKey)
                                           .setSalt(salt)
                                           .buildForEncryptingByPassword();

        EncryptionUtils.encryptFile(encParam, sourceFile, encryptedFile);

        checkEncryptedFile();

        EncryptionParameter decParam =
                EncryptionParameter.Builder.create()
                                           .setAlgorithm(ALGORITHM_AES)
                                           .setStrength(STRENGTH_256)
                                           .setPassword(masterKey)
                                           .setSalt(encParam.getSalt())
                                           .setInitVector(encParam.getInitVector())
                                           .buildForDecryptingByPassword();

        EncryptionUtils.decryptFile(decParam, encryptedFile, decryptedFile);

        checkDecryptedFile();
    }

    private byte[] genSalt() throws NoSuchAlgorithmException {
        SecureRandom random = SecureRandom.getInstance(EncryptionParameter.RANDOM_ALGORITHM);

        byte[] salt = new byte[8];
        random.nextBytes(salt);

        return salt;
    }
}

說明

我們在加密時，需要一些參數，有些參數是從外部取得，例如使用者輸入(masterKey)，或者是亂數生成(salt)。有些參數雖然不是由我們自己生成，但是需要被保留。

這裡有兩個測試，一個是用最簡單的參數，直接進行加解密的操作；另一個則是可以接受一個外部輸入的密碼，以及事先預備好的鹽巴(salt)，以便產生加密金鑰(SecretKey)。

然而，這個程式除了透露出我不太會寫 TestCase，還透露出我們準備加解密用的參數，需要哪些內容。例如在第一個測試裡，我們給了最簡單的兩個參數，指定演算法是 AES 以及金鑰長度是 256；另一個測試，則是增加了密碼與鹽巴。在這兩個加密參數賦值之後，有兩個不一樣的 build 名稱，有別於一般 Builder 模式所使用的。這是因為我們建構的目的不一樣，因此名稱也要不同，以表明內部執行方式的不同。

EncryptionParameter

package javaxx.cipher;

import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.util.HashMap;
import java.util.Map;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringUtils;

/**
 * EncryptionParameter
 * cchuang, 2016/12/4.
 */
public class EncryptionParameter {

    public static final String ALGORITHM_AES = "AES";
    public static final int STRENGTH_256 = 256;
    public static final String AES_IMPLEMENT_DEFAULT = "AES/CTR/PKCS5Padding";
    public static final String RANDOM_ALGORITHM = "SHA1PRNG";
    public static final String SECRET_KEY_ALGORITHM = "PBKDF2WithHmacSHA256";
    public static final int INTERATION_COUNT = 65536;

    public static final Map<String, String> ALGORITHMS = new HashMap<>();

    static {
        ALGORITHMS.put(ALGORITHM_AES, AES_IMPLEMENT_DEFAULT);
    }

    private String algorithmImpl;
    private byte[] initVector;
    private byte[] salt;
    private SecretKey secretKey;

    private EncryptionParameter() {}

    public EncryptionParameter(String algorithm, int strength) throws NoSuchAlgorithmException {
        this.algorithmImpl = ALGORITHMS.get(algorithm);
        this.secretKey = initSecretKey(algorithm, strength);
        this.initVector = initInitialVector(strength);
    }

    private EncryptionParameter(String algorithm, SecretKey secretKey, byte[] iv) {
        this.algorithmImpl = ALGORITHMS.get(algorithm);
        this.secretKey = secretKey;
        this.initVector = iv;
    }

    private EncryptionParameter(String algorithm, int strength, String password, byte[] salt)
            throws InvalidKeySpecException, NoSuchAlgorithmException {
        this.algorithmImpl = ALGORITHMS.get(algorithm);
        this.salt = salt;
        this.secretKey = initSecretKeyByPassword(algorithm, strength, password, salt);
        this.initVector = initInitialVector(strength);
    }

    private SecretKey initSecretKeyByPassword(String algorithm, int strength, String password, byte[] salt)
            throws NoSuchAlgorithmException, InvalidKeySpecException {
        SecretKey secretKey = generateSecretKeyByCostomPasswordAndSalt(password, salt, algorithm, strength);

        return secretKey;
    }

    public EncryptionParameter(String algorithm, int strength, String password, byte[] salt, byte[] iv)
            throws InvalidKeySpecException, NoSuchAlgorithmException {
        this.algorithmImpl = ALGORITHMS.get(algorithm);
        this.secretKey = initSecretKeyByPassword(algorithm, strength, password, salt);
        this.initVector = iv;
    }

    private SecretKey generateSecretKeyByCostomPasswordAndSalt(String password, byte[] salt, String algorithm,
            int strength) throws InvalidKeySpecException, NoSuchAlgorithmException {
        SecretKeyFactory secKeyFactory = SecretKeyFactory.getInstance(SECRET_KEY_ALGORITHM);

        KeySpec spec = new PBEKeySpec(password.toCharArray(), salt, INTERATION_COUNT, strength);
        SecretKey temp = secKeyFactory.generateSecret(spec);

        return new SecretKeySpec(temp.getEncoded(), algorithm);
    }

    private SecretKey initSecretKey(String algorithm, int strength) throws NoSuchAlgorithmException {
        KeyGenerator keyGen = KeyGenerator.getInstance(algorithm);

        keyGen.init(strength, SecureRandom.getInstance(RANDOM_ALGORITHM));

        SecretKey secretKey = keyGen.generateKey();

        return secretKey;
    }

    private byte[] initInitialVector(int strength) throws NoSuchAlgorithmException {
        byte[] initVector = new byte[128 / 8];

        SecureRandom random = SecureRandom.getInstance(RANDOM_ALGORITHM);
        random.nextBytes(initVector);

        return initVector;
    }

    public byte[] getInitVector() {
        return initVector;
    }

    public SecretKey getSecretKey() {
        return secretKey;
    }

    public String getAlgorithmImpl() {
        return algorithmImpl;
    }

    public byte[] getSalt() {
        return salt;
    }

    public static final class Builder {
        private String algorithm;
        private byte[] initVector;
        private byte[] salt;
        private SecretKey secretKey;
        private int strength = -1;
        private String password;

        private Builder() {}

        public static Builder create() { return new Builder();}

        public Builder setAlgorithm(String algorithm) {
            this.algorithm = algorithm;
            return this;
        }

        public Builder setStrength(int strength) throws NoSuchAlgorithmException {
            this.strength = strength;
            return this;
        }

        public Builder setInitVector(byte[] initVector) {
            this.initVector = initVector;
            return this;
        }

        public Builder setSalt(byte[] salt) {
            this.salt = salt;
            return this;
        }

        public Builder setSecretKey(SecretKey secretKey) {
            this.secretKey = secretKey;
            return this;
        }

        public Builder setPassword(String password) {
            this.password = password;
            return this;
        }

        public EncryptionParameter buildForDirectlyInitial() {
            if (!isDirectlyInitial()) {
                throw new IllegalArgumentException();
            }
            return new EncryptionParameter(algorithm, secretKey, initVector);
        }

        private boolean isDirectlyInitial() {
            return !StringUtils.isEmpty(algorithm)
                    && null != secretKey
                    && !ArrayUtils.isEmpty(initVector);
        }

        public EncryptionParameter buildForDecryptingByPassword()
                throws InvalidKeySpecException, NoSuchAlgorithmException {
            if (!isDecryptingByPassword()) {
                throw new IllegalArgumentException();
            }
            return new EncryptionParameter(algorithm, strength, password, salt, initVector);
        }

        private boolean isDecryptingByPassword() {
            return !StringUtils.isEmpty(algorithm)
                    && 0 < strength
                    && !StringUtils.isEmpty(password)
                    && !ArrayUtils.isEmpty(initVector)
                    && !ArrayUtils.isEmpty(salt);
        }

        public EncryptionParameter buildForEncryptingByPassword()
                throws InvalidKeySpecException, NoSuchAlgorithmException {
            if (!isEncryptingByPassword()) {
                throw new IllegalArgumentException();
            }
            return new EncryptionParameter(algorithm, strength, password, salt);
        }

        private boolean isEncryptingByPassword() {
            return !StringUtils.isEmpty(algorithm)
                    && 0 < strength
                    && !StringUtils.isEmpty(password)
                    && !ArrayUtils.isEmpty(salt);
        }

        public EncryptionParameter buildForEncryption() throws NoSuchAlgorithmException {
            if (!isDefaultEncryption()) {
                throw new IllegalArgumentException();
            }
            return new EncryptionParameter(algorithm, strength);
        }

        private boolean isDefaultEncryption() {
            return !StringUtils.isEmpty(algorithm)
                    && 0 < strength;
        }
    }
}

這個類別，早期的版本是用 getInstance(...) 的方式來取得，但這樣表達並不完整，而且語意並不明白。因此改為 Builder pattern，讓閱讀的人，知道需要預備什麼參數，以及要以何種意圖來呼叫它。

加密的實作方式，我們使用 AES/CTR/PKCS5Padding，但因為我不是研究密碼學的，沒有辦法詳細說明這裡面的每一個字眼，但現行 AES256 加密的 ECB 實作方式，不再建議使用。

EncryptionUtils

package javaxx.cipher;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;

/**
 * EncryptionUtils
 * cchuang, 2016/12/4.
 */
public class EncryptionUtils {

    public static void encryptFile(EncryptionParameter param, File source, File encrypted)
            throws NoSuchPaddingException, NoSuchAlgorithmException, InvalidKeyException,
            InvalidAlgorithmParameterException, IOException, BadPaddingException, IllegalBlockSizeException {
        if (null == param || null == source || null == encrypted || !source.exists()) {
            throw new RuntimeException();
        }

        Cipher cipher = initEncryptedCipher(param);

        readSourceAndProcess(cipher, source, encrypted);
    }

    private static void readSourceAndProcess(Cipher cipher, File source, File target)
            throws BadPaddingException, IllegalBlockSizeException, IOException {
        byte[] buffer = new byte[4096];

        try (InputStream fis = new FileInputStream(source);
                OutputStream fos = new FileOutputStream(target)) {

            int readBytes = fis.read(buffer);

            while (readBytes > 0) {
                byte[] cipherBytes = cipher.update(buffer, 0, readBytes);
                fos.write(cipherBytes);
                readBytes = fis.read(buffer);
            }

            cipher.doFinal();
        }
    }

    private static Cipher initEncryptedCipher(EncryptionParameter param)
            throws InvalidAlgorithmParameterException, InvalidKeyException, NoSuchPaddingException,
            NoSuchAlgorithmException {
        Cipher cipher = Cipher.getInstance(param.getAlgorithmImpl());

        cipher.init(Cipher.ENCRYPT_MODE, param.getSecretKey(), new IvParameterSpec(param.getInitVector()));

        return cipher;
    }

    public static void decryptFile(EncryptionParameter param, File encrypted, File decrypted)
            throws BadPaddingException, IllegalBlockSizeException, IOException, NoSuchAlgorithmException,
            NoSuchPaddingException, InvalidAlgorithmParameterException, InvalidKeyException {
        if (null == param || null == encrypted || null == decrypted || !encrypted.exists()) {
            throw new RuntimeException();
        }

        Cipher cipher = initDecryptedCipher(param);

        if (!decrypted.exists()) {
            decrypted.createNewFile();
        }

        readSourceAndProcess(cipher, encrypted, decrypted);
    }

    private static Cipher initDecryptedCipher(EncryptionParameter param)
            throws InvalidAlgorithmParameterException, InvalidKeyException, NoSuchPaddingException,
            NoSuchAlgorithmException {
        Cipher cipher = Cipher.getInstance(param.getAlgorithmImpl());

        cipher.init(Cipher.DECRYPT_MODE, param.getSecretKey(), new IvParameterSpec(param.getInitVector()));

        return cipher;
    }
}

最後，這是我們的主程式，依據傳來的 parameter 進行加密或解密。

小結

如同第二個測試所示，執行完後，我們會有 iv 與 salt 要帶到解密時期。這時我們可以選擇把這兩個值附在檔案的頭或尾，或轉成 Base64 給對方。或者，我們慢慢地擴展這個作法。讓它可以使用 RSA 對一些外部資訊，諸如 masterKey、iv、salt 進行加密。