package com.vsoontech.p2p.sample;

import org.bouncycastle.jce.provider.BouncyCastleProvider;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.security.Key;
import java.security.NoSuchAlgorithmException;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

/**
 * @author zhou
 * @since 2016/9/26
 */

public class AESKeyModel {
    public static final String KEY_ALGORITHM = "AES";
    private static final String DEFAULT_CIPHER_ALGORITHM = "AES/CBC/PKCS5Padding";
    private String srcFile = "", destionFile = "";

    /**
     * 初始化金鑰
     *
     * @return byte[] 金鑰
     * @throws Exception
     */
    public byte[] initSecretKey() {
        //返回生成指定演算法的祕密金鑰的 KeyGenerator 物件
        KeyGenerator kg = null;
        try {
            kg = KeyGenerator.getInstance(KEY_ALGORITHM);
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
            return new byte[0];
        }
        //初始化此金鑰生成器，使其具有確定的金鑰大小
        //AES 要求金鑰長度為 128
        kg.init(128);
        //生成一個金鑰
        SecretKey secretKey = kg.generateKey();
        return secretKey.getEncoded();
    }

    public void setDestionFile(String destionFile) {
        this.destionFile = destionFile;
    }

    public void setSrcFile(String srcFile) {
        this.srcFile = srcFile;
    }

    /**
     * 轉換金鑰
     *
     * @param key 二進位制金鑰
     * @return 金鑰
     */
    private static Key toKey(byte[] key) {
        //生成金鑰
        return new SecretKeySpec(key, KEY_ALGORITHM);
    }

    /**
     * 加密
     *
     * @param data 待加密資料
     * @param key  金鑰
     * @return byte[]   加密資料
     * @throws Exception
     */
    public static byte[] encrypt(byte[] data, Key key) throws Exception {
        return encrypt(data, key, DEFAULT_CIPHER_ALGORITHM);
    }

    /**
     * 加密
     *
     * @param data 待加密資料
     * @param key  二進位制金鑰
     * @return byte[]   加密資料
     * @throws Exception
     */
    public static byte[] encrypt(byte[] data, byte[] key) throws Exception {
        return encrypt(data, key, DEFAULT_CIPHER_ALGORITHM);
    }


    /**
     * 加密
     *
     * @param data            待加密資料
     * @param key             二進位制金鑰
     * @param cipherAlgorithm 加密演算法/工作模式/填充方式
     * @return byte[]   加密資料
     * @throws Exception
     */
    public static byte[] encrypt(byte[] data, byte[] key, String cipherAlgorithm) throws Exception {
        //還原金鑰
        Key k = toKey(key);
        return encrypt(data, k, cipherAlgorithm);
    }

    /**
     * 加密
     *
     * @param data            待加密資料
     * @param key             金鑰
     * @param cipherAlgorithm 加密演算法/工作模式/填充方式
     * @return byte[]   加密資料
     * @throws Exception
     */
    public static byte[] encrypt(byte[] data, Key key, String cipherAlgorithm) throws Exception {
        //例項化
        Cipher cipher = Cipher.getInstance(cipherAlgorithm);
        //使用金鑰初始化，設定為加密模式
        cipher.init(Cipher.ENCRYPT_MODE, key);
        //執行操作
        return cipher.doFinal(data);
    }

    /**
     * 解密
     *
     * @param data 待解密資料
     * @param key  二進位制金鑰
     * @return byte[]   解密資料
     * @throws Exception
     */
    public static byte[] decrypt(byte[] data, byte[] key) throws Exception {
        return decrypt(data, key, DEFAULT_CIPHER_ALGORITHM);
    }

    /**
     * 解密
     *
     * @param data 待解密資料
     * @param key  金鑰
     * @return byte[]   解密資料
     * @throws Exception
     */
    public static byte[] decrypt(byte[] data, Key key) throws Exception {
        return decrypt(data, key, DEFAULT_CIPHER_ALGORITHM);
    }

    /**
     * 解密
     *
     * @param data            待解密資料
     * @param key             二進位制金鑰
     * @param cipherAlgorithm 加密演算法/工作模式/填充方式
     * @return byte[]   解密資料
     * @throws Exception
     */
    public static byte[] decrypt(byte[] data, byte[] key, String cipherAlgorithm) throws Exception {
        //還原金鑰
        Key k = toKey(key);
        return decrypt(data, k, cipherAlgorithm);
    }

    /**
     * 解密
     *
     * @param data            待解密資料
     * @param key             金鑰
     * @param cipherAlgorithm 加密演算法/工作模式/填充方式
     * @return byte[]   解密資料
     * @throws Exception
     */
    public static byte[] decrypt(byte[] data, Key key, String cipherAlgorithm) throws Exception {
        //例項化
        Cipher cipher = Cipher.getInstance(cipherAlgorithm);
        //使用金鑰初始化，設定為解密模式
        cipher.init(Cipher.DECRYPT_MODE, key);
        //執行操作
        return cipher.doFinal(data);
    }

    public void encryptionFile(Key sessionKey) throws Exception {
        int len = 0;
        byte[] buffer = new byte[1024];
        byte[] cipherbuffer = null;

        // 使用會話金鑰對檔案加密。
        Cipher cipher = Cipher.getInstance(DEFAULT_CIPHER_ALGORITHM, new BouncyCastleProvider());
        IvParameterSpec iv = new IvParameterSpec("0000000000123456".getBytes());
        cipher.init(Cipher.ENCRYPT_MODE, sessionKey, iv);

        FileInputStream fis = new FileInputStream(new File(srcFile));
        FileOutputStream fos = new FileOutputStream(new File(destionFile));

        // 讀取原文，加密並寫密文到輸出檔案。
        while ((len = fis.read(buffer)) != -1) {
            cipherbuffer = cipher.update(buffer, 0, len);
            fos.write(cipherbuffer);
            fos.flush();
        }
        cipherbuffer = cipher.doFinal();
        fos.write(cipherbuffer);
        fos.flush();

        if (fis != null)
            fis.close();
        if (fos != null)
            fos.close();
    }

    public void descryptionFile(Key sessionKey) throws Exception {
        int len = 0;
        byte[] buffer = new byte[5 * 1024];
        byte[] plainbuffer = null;

        Cipher cipher = Cipher.getInstance(DEFAULT_CIPHER_ALGORITHM, new BouncyCastleProvider());
        IvParameterSpec iv = new IvParameterSpec("0000000000123456".getBytes());
        cipher.init(Cipher.DECRYPT_MODE, sessionKey, iv);

        FileInputStream fis = new FileInputStream(new File(srcFile));
        FileOutputStream fos = new FileOutputStream(new File(destionFile));

        while ((len = fis.read(buffer)) != -1) {
            plainbuffer = cipher.update(buffer, 0, len);
            fos.write(plainbuffer);
            fos.flush();
        }

        plainbuffer = cipher.doFinal();
        fos.write(plainbuffer);
        fos.flush();

        if (fis != null)
            fis.close();
        if (fos != null)
            fos.close();
    }
}