package com.palacesun.masking.tools.thkms.util.gmhelper;

import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.SecureRandom;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.Mac;
import org.bouncycastle.crypto.engines.SM4Engine;
import org.bouncycastle.crypto.macs.CBCBlockCipherMac;
import org.bouncycastle.crypto.macs.GMac;
import org.bouncycastle.crypto.modes.GCMBlockCipher;
import org.bouncycastle.crypto.paddings.BlockCipherPadding;
import org.bouncycastle.crypto.paddings.PKCS7Padding;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;

public class SM4Util extends GMBaseUtil {
   public static final String ALGORITHM_NAME = "SM4";
   public static final String ALGORITHM_NAME_ECB_PADDING = "SM4/ECB/PKCS5Padding";
   public static final String ALGORITHM_NAME_ECB_NOPADDING = "SM4/ECB/NoPadding";
   public static final String ALGORITHM_NAME_CBC_PADDING = "SM4/CBC/PKCS5Padding";
   public static final String ALGORITHM_NAME_CBC_NOPADDING = "SM4/CBC/NoPadding";
   public static final int DEFAULT_KEY_SIZE = 128;

   public static byte[] generateKey() throws NoSuchAlgorithmException, NoSuchProviderException {
      return generateKey(128);
   }

   public static byte[] generateKey(int keySize) throws NoSuchAlgorithmException, NoSuchProviderException {
      KeyGenerator kg = KeyGenerator.getInstance("SM4", "BC");
      kg.init(keySize, new SecureRandom());
      return kg.generateKey().getEncoded();
   }

   public static byte[] encrypt_ECB_Padding(byte[] key, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException {
      Cipher cipher = generateECBCipher("SM4/ECB/PKCS5Padding", 1, key);
      return cipher.doFinal(data);
   }

   public static byte[] decrypt_ECB_Padding(byte[] key, byte[] cipherText) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException {
      Cipher cipher = generateECBCipher("SM4/ECB/PKCS5Padding", 2, key);
      return cipher.doFinal(cipherText);
   }

   public static byte[] encrypt_ECB_NoPadding(byte[] key, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException {
      Cipher cipher = generateECBCipher("SM4/ECB/NoPadding", 1, key);
      return cipher.doFinal(data);
   }

   public static byte[] decrypt_ECB_NoPadding(byte[] key, byte[] cipherText) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException {
      Cipher cipher = generateECBCipher("SM4/ECB/NoPadding", 2, key);
      return cipher.doFinal(cipherText);
   }

   public static byte[] encrypt_CBC_Padding(byte[] key, byte[] iv, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
      Cipher cipher = generateCBCCipher("SM4/CBC/PKCS5Padding", 1, key, iv);
      return cipher.doFinal(data);
   }

   public static byte[] decrypt_CBC_Padding(byte[] key, byte[] iv, byte[] cipherText) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidAlgorithmParameterException {
      Cipher cipher = generateCBCCipher("SM4/CBC/PKCS5Padding", 2, key, iv);
      return cipher.doFinal(cipherText);
   }

   public static byte[] encrypt_CBC_NoPadding(byte[] key, byte[] iv, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
      Cipher cipher = generateCBCCipher("SM4/CBC/NoPadding", 1, key, iv);
      return cipher.doFinal(data);
   }

   public static byte[] decrypt_CBC_NoPadding(byte[] key, byte[] iv, byte[] cipherText) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidAlgorithmParameterException {
      Cipher cipher = generateCBCCipher("SM4/CBC/NoPadding", 2, key, iv);
      return cipher.doFinal(cipherText);
   }

   public static byte[] doCMac(byte[] key, byte[] data) throws NoSuchProviderException, NoSuchAlgorithmException, InvalidKeyException {
      Key keyObj = new SecretKeySpec(key, "SM4");
      return doMac("SM4-CMAC", keyObj, data);
   }

   public static byte[] doGMac(byte[] key, byte[] iv, int tagLength, byte[] data) {
      Mac mac = new GMac(new GCMBlockCipher(new SM4Engine()), tagLength * 8);
      return doMac(mac, key, iv, data);
   }

   public static byte[] doCBCMac(byte[] key, byte[] iv, byte[] data) {
      SM4Engine engine = new SM4Engine();
      Mac mac = new CBCBlockCipherMac(engine, engine.getBlockSize() * 8, new PKCS7Padding());
      return doMac(mac, key, iv, data);
   }

   public static byte[] doCBCMac(byte[] key, byte[] iv, BlockCipherPadding padding, byte[] data) throws Exception {
      SM4Engine engine = new SM4Engine();
      if (padding == null && data.length % engine.getBlockSize() != 0) {
         throw new Exception("if no padding, data length must be multiple of SM4 BlockSize");
      } else {
         Mac mac = new CBCBlockCipherMac(engine, engine.getBlockSize() * 8, padding);
         return doMac(mac, key, iv, data);
      }
   }

   private static byte[] doMac(Mac mac, byte[] key, byte[] iv, byte[] data) {
      CipherParameters cipherParameters = new KeyParameter(key);
      mac.init(new ParametersWithIV(cipherParameters, iv));
      mac.update(data, 0, data.length);
      byte[] result = new byte[mac.getMacSize()];
      mac.doFinal(result, 0);
      return result;
   }

   private static byte[] doMac(String algorithmName, Key key, byte[] data) throws NoSuchProviderException, NoSuchAlgorithmException, InvalidKeyException {
      javax.crypto.Mac mac = javax.crypto.Mac.getInstance(algorithmName, "BC");
      mac.init(key);
      mac.update(data);
      return mac.doFinal();
   }

   private static Cipher generateECBCipher(String algorithmName, int mode, byte[] key) throws NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidKeyException {
      Cipher cipher = Cipher.getInstance(algorithmName, "BC");
      Key sm4Key = new SecretKeySpec(key, "SM4");
      cipher.init(mode, sm4Key);
      return cipher;
   }

   private static Cipher generateCBCCipher(String algorithmName, int mode, byte[] key, byte[] iv) throws InvalidKeyException, InvalidAlgorithmParameterException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException {
      Cipher cipher = Cipher.getInstance(algorithmName, "BC");
      Key sm4Key = new SecretKeySpec(key, "SM4");
      IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
      cipher.init(mode, sm4Key, ivParameterSpec);
      return cipher;
   }
}
