elink-online-common/src/main/java/com/eactive/eai/common/security/AESCryptoModuleExtension.java

package com.eactive.eai.common.security;

import java.nio.ByteBuffer;
import java.security.Key;
import java.security.SecureRandom;
import java.security.Security;
import java.security.spec.AlgorithmParameterSpec;

import javax.crypto.Cipher;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.bouncycastle.jcajce.spec.FPEParameterSpec;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

public class AESCryptoModuleExtension implements CryptoModuleExtension {

    private static final int    GCM_IV_LENGTH     = 12;
    private static final int    GCM_TAG_LENGTH    = 128;
    private static final String DEFAULT_PROVIDER  = "BC";

    private Cipher encryptEngine;
    private Cipher decryptEngine;
    private String mode;
    private String pad;
    private byte[] iv;
    private byte[] encKey;
    private byte[] decKey;
    /** 명시적 프로바이더. null이면 모든 모드에서 BC(DEFAULT_PROVIDER)를 사용한다. */
    private String provider;

    // ------------------------------------------------------------------
    // init — 모든 오버로드는 최종적으로 6+provider 버전에 위임한다
    // ------------------------------------------------------------------

    @Override
    public void init(String alg, String mode, String pad, byte[] encKey, byte[] decKey) throws Exception {
        init(alg, mode, pad, null, encKey, decKey, null);
    }

    @Override
    public void init(String alg, String mode, String pad, byte[] iv, byte[] encKey, byte[] decKey) throws Exception {
        init(alg, mode, pad, iv, encKey, decKey, null);
    }

    @Override
    public void init(String alg, String mode, String pad, byte[] encKey, byte[] decKey, String provider) throws Exception {
        init(alg, mode, pad, null, encKey, decKey, provider);
    }

    @Override
    public void init(String alg, String mode, String pad, byte[] iv, byte[] encKey, byte[] decKey, String provider) throws Exception {
        this.mode     = mode;
        this.pad      = pad;
        this.iv       = iv;
        this.encKey   = encKey;
        this.decKey   = decKey;
        this.provider = (provider != null && !provider.isEmpty()) ? provider : null;

        // GCM 포함 모든 모드에서 BC 프로바이더를 사용 — 기본값 BC
        if (Security.getProvider("BC") == null) {
            Security.addProvider(new BouncyCastleProvider());
        }

        if (!mode.equalsIgnoreCase("GCM")) {
            String m = mode != null ? mode : "";
            String p = pad != null ? pad : "";
            String transformation = String.format("AES/%s/%s", m, p);

            Key encKeySpec = new SecretKeySpec(encKey, "AES");
            Key decKeySpec = new SecretKeySpec(decKey, "AES");
            AlgorithmParameterSpec param = iv != null ? new IvParameterSpec(iv) : null;
            if (m.toLowerCase().contains("ff")) {
                param = new FPEParameterSpec(256, iv);
            }

            String prov = this.provider != null ? this.provider : DEFAULT_PROVIDER;
            this.encryptEngine = Cipher.getInstance(transformation, prov);
            this.decryptEngine = Cipher.getInstance(transformation, prov);
            this.encryptEngine.init(Cipher.ENCRYPT_MODE, encKeySpec, param, (SecureRandom) null);
            this.decryptEngine.init(Cipher.DECRYPT_MODE, decKeySpec, param, (SecureRandom) null);
        }
    }

    // ------------------------------------------------------------------
    // encrypt
    // ------------------------------------------------------------------

    @Override
    public int encrypt(byte[] src, int sindex, int slength, byte[] dst, int dindex) throws Exception {
        return this.encryptEngine.doFinal(src, sindex, slength, dst, dindex);
    }

    @Override
    public byte[] encrypt(byte[] src, int sindex, int slength) throws Exception {
        return encrypt(src, sindex, slength, null);
    }

    /**
     * AAD를 명시한 암호화.
     * GCM 모드에서 aad != null 이면 해당 값을 사용하고, null 이면 iv를 AAD로 사용하는 기본 동작을 따른다.
     * 비-GCM 모드는 AAD를 무시하고 사전 초기화된 엔진으로 처리한다.
     */
    @Override
    public byte[] encrypt(byte[] src, int sindex, int slength, byte[] aad) throws Exception {
        if (this.mode.equalsIgnoreCase("GCM")) {
            Key key = new SecretKeySpec(this.encKey, "AES");
            byte[] nonce = new byte[GCM_IV_LENGTH];
            new SecureRandom().nextBytes(nonce);
            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH, nonce);

            String prov = this.provider != null ? this.provider : DEFAULT_PROVIDER;
            Cipher cipher = Cipher.getInstance(String.format("AES/%s/%s", this.mode, this.pad), prov);
            cipher.init(Cipher.ENCRYPT_MODE, key, gcmSpec);

            byte[] effectiveAad = aad != null ? aad : this.iv;
            if (effectiveAad != null) cipher.updateAAD(effectiveAad);

            byte[] cipherTextWithTag = cipher.doFinal(src, sindex, slength);
            byte[] output = new byte[GCM_IV_LENGTH + cipherTextWithTag.length];
            System.arraycopy(nonce, 0, output, 0, GCM_IV_LENGTH);
            System.arraycopy(cipherTextWithTag, 0, output, GCM_IV_LENGTH, cipherTextWithTag.length);
            return output;
        }

        int outLen = this.encryptEngine.getOutputSize(slength);
        byte[] dst  = new byte[outLen];
        int actual  = encrypt(src, sindex, slength, dst, 0);
        return ByteBuffer.allocate(actual).put(dst, 0, actual).array();
    }

    // ------------------------------------------------------------------
    // decrypt
    // ------------------------------------------------------------------

    @Override
    public int decrypt(byte[] src, int sindex, int slength, byte[] dst, int dindex) throws Exception {
        return this.decryptEngine.doFinal(src, sindex, slength, dst, dindex);
    }

    @Override
    public byte[] decrypt(byte[] src, int sindex, int slength) throws Exception {
        return decrypt(src, sindex, slength, null);
    }

    /**
     * AAD를 명시한 복호화.
     * GCM 모드에서 aad != null 이면 해당 값을 사용하고, null 이면 iv를 AAD로 사용하는 기본 동작을 따른다.
     * 비-GCM 모드는 AAD를 무시하고 사전 초기화된 엔진으로 처리한다.
     */
    @Override
    public byte[] decrypt(byte[] src, int sindex, int slength, byte[] aad) throws Exception {
        if (this.mode.equalsIgnoreCase("GCM")) {
            Key key = new SecretKeySpec(this.decKey, "AES");
            byte[] nonce = new byte[GCM_IV_LENGTH];
            System.arraycopy(src, sindex, nonce, 0, GCM_IV_LENGTH);
            byte[] cipherTextWithTag = new byte[slength - GCM_IV_LENGTH];
            System.arraycopy(src, sindex + GCM_IV_LENGTH, cipherTextWithTag, 0, cipherTextWithTag.length);

            GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH, nonce);
            String prov = this.provider != null ? this.provider : DEFAULT_PROVIDER;
            Cipher cipher = Cipher.getInstance(String.format("AES/%s/%s", this.mode, this.pad), prov);
            cipher.init(Cipher.DECRYPT_MODE, key, gcmSpec);

            byte[] effectiveAad = aad != null ? aad : this.iv;
            if (effectiveAad != null) cipher.updateAAD(effectiveAad);

            return cipher.doFinal(cipherTextWithTag);
        }
        return this.decryptEngine.doFinal(src, sindex, slength);
    }
}