Flink实战- （8）Flink UDF实现数据脱敏

背景

        在Flink环境实现数据实时同步的过程中，需要对某些字段进行脱敏处理或者特殊处理，这是需要自己定义UDF函数进行转换。

代码实现

1. AES加密UDF

import javax.crypto.Cipher;
import javax.crypto.spec.SecretKeySpec;
import org.apache.flink.table.functions.ScalarFunction;

public class AesEncryptUDF extends ScalarFunction {
    private static final byte[] KEY = "0123456789abcdef".getBytes(); // AES 密钥
    private static final String CIPHER_ALGORITHM = "AES/ECB/PKCS5Padding"; // 加密算法

    public String eval(String content) {
        try {
            SecretKeySpec keySpec = new SecretKeySpec(KEY, "AES");
            Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, keySpec);
            byte[] encrypted = cipher.doFinal(content.getBytes());
            return new String(encrypted);
        } catch (Exception e) {
            return null;
        }
    }
}

使用该 UDF 可以在 Flink SQL 中对某个字段的值进行 AES 加密，示例代码如下：

SELECT AesEncryptUDF(column_name) as encrypted_column FROM table_name

2. MD5加密UDF

import java.security.MessageDigest;
import org.apache.flink.table.functions.ScalarFunction;

public class Md5EncryptUDF extends ScalarFunction {
    public String eval(String content) {
        try {
            MessageDigest messageDigest = MessageDigest.getInstance("MD5");
            byte[] md5Bytes = messageDigest.digest(content.getBytes());
            StringBuilder hexString = new StringBuilder();
            for (byte md5Byte : md5Bytes) {
                String hex = Integer.toHexString(0xff & md5Byte);
                if (hex.length() == 1) {
                    hexString.append('0');
                }
                hexString.append(hex);
            }
            return hexString.toString();
        } catch (Exception e) {
            return null;
        }
    }
}

使用该 UDF 可以在 Flink SQL 中对某个字段的值进行 MD5 加密，示例代码如下：

SELECT Md5EncryptUDF(column_name) as encrypted_column FROM table_name

3.SM4加密UDF

import org.apache.commons.codec.binary.Hex;
import org.apache.commons.lang.StringUtils;
import org.bouncycastle.crypto.CipherKeyGenerator;
import org.bouncycastle.crypto.engines.SM4Engine;
import org.bouncycastle.crypto.generators.SCrypt;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ScryptParameters;
import org.bouncycastle.util.encoders.Base64;
import org.bouncycastle.util.encoders.HexEncoder;
import org.apache.flink.table.functions.ScalarFunction;

public class Sm4EncryptUDF extends ScalarFunction {
    //SM4 密钥
    private static final byte[] KEY = "0123456789abcdef".getBytes();
    //SM4 密码学模式
    private static final String CIPHER_ALGORITHM = "SM4/ECB/PKCS7Padding";
    //SM4 密钥生成器
    private static final CipherKeyGenerator KEY_GENERATOR = new CipherKeyGenerator();
    //SM4 参数
    private static final ScryptParameters SCRYPT_PARAMETERS = new ScryptParameters(16384, 8, 1, 32);

    public String eval(String content) {
        if (StringUtils.isBlank(content)) {
            return null;
        }

        try {
            byte[] salt = new byte[16];
            KEY_GENERATOR.init(SCRYPT_PARAMETERS);
            byte[] sm4Key = new byte[16];
            KEY_GENERATOR.generateKey(sm4Key, 0, 16);
            KeyParameter keyParameter = new KeyParameter(sm4Key);
            SM4Engine engine = new SM4Engine();
            engine.init(true, keyParameter);
            byte[] input = content.getBytes();

            int cipherLen = engine.getOutputSize(input.length);
            byte[] cipherBytes = new byte[cipherLen];
            int cipherLength = engine.processBytes(input, 0, input.length, cipherBytes, 0);
            engine.doFinal(cipherBytes, cipherLength);

            HexEncoder hex = new HexEncoder();
            String cipherText = new String(Base64.encode(cipherBytes), "UTF-8");
            return cipherText;
        } catch (Exception e) {
            return null;
        }
    }
}

使用该 UDF 可以在 Flink SQL 中对某个字段的值进行 SM4 加密，示例代码如下：

SELECT Sm4EncryptUDF(column_name) as encrypted_column FROM table_name

4. RSA加密UDF

import java.nio.charset.StandardCharsets;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;

import javax.crypto.Cipher;

import org.apache.flink.table.functions.ScalarFunction;

public class RsaEncryptUDF extends ScalarFunction {

    public String eval(String content) {
        try {
            // 生成 RSA 密钥对
            KeyPair kp = generateRSAKeyPair();

            // 获取公钥和私钥
            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.ENCRYPT_MODE, kp.getPublic());
            byte[] encryptedBytes = cipher.doFinal(content.getBytes(StandardCharsets.UTF_8));

            return Base64.encodeBase64String(encryptedBytes);
        } catch (Exception e) {
            return null;
        }
    }

    private static KeyPair generateRSAKeyPair() throws NoSuchAlgorithmException {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048); // 生成 2048 位密钥对
        return keyPairGenerator.genKeyPair();
    }
}

使用该 UDF 可以在 Flink SQL 中对某个字段的值进行 RSA 加密，示例代码如下：

SELECT RsaEncryptUDF(column_name) as encrypted_column FROM table_name