Spark 实现自定义加密

Spark 实现自定义加密

一、建立加密和解密的自定义函数

import java.nio.charset.{StandardCharsets}
import java.util.Base64
import javax.crypto.Cipher
import javax.crypto.spec.SecretKeySpec

object SparkUtil {
	 /**
   * 处理密钥
   * @param secret 密钥
   */
  private def secretInit(secret:String)={
    // 对密钥长度进行约束
    val allowNumBits: Array[Int] = Array(16, 24, 32)
    // 如果密钥长度符合，将密钥转换为AES密钥对象
    if (allowNumBits.contains(secret.size)) {
      new SecretKeySpec(
          secret.getBytes(StandardCharsets.UTF_8),"AES")
    }else{
      throw new RuntimeException(
          s"AES secret size of numBits ${secret.size} not in 
          permitted values (${allowNumBits.mkString(",")})")
    }
  }
  /**
   * 加密函数
   * @param src 源数据
   * @param secret 密钥
   */
  def encrypt(src:String,secret:String)={
    // 获取加密算法实例
    val cipher: Cipher = Cipher.getInstance("AES")
    // 初始化加密模式，使用给定的密钥(需要先用key()对密钥进行处理)
    cipher.init(Cipher.ENCRYPT_MODE,secretInit(secret))
    // 执行加密操作
    val bytes: Array[Byte] = cipher.doFinal(src.getBytes(StandardCharset.UTF_8))
    // 返回加密后的数据
    Base64.getEncoder().encodeToString(bytes)
  }
  /**
   * 解密函数
   * @param dest 待解密数据
   * @param secret 密钥
   */
  def decrypt(dest:String,secret:String)={
    val cipher: Cipher = Cipher.getInstance("AES")
    cipher.init(Cipher.DECRYPT_MODE,secretInit(secret))
    val bytes: Array[Byte] = cipher.doFinal(
        Base64.getDecoder.decode(dest))
    new String(bytes, StandardCharsets.UTF_8)
  }
}

二、在 Spark 环境下导入对象实现的方法，并在 SparkSession 中注册 UDF 函数

import core.SparkUtil.{encrypt,decrypt}
spark.udf.register(
    "aes_encrypt",
    (src:String,secret:String) 
    	=>encrypt(src, secret),StringType)
spark.udf.register(
    "aes_decrypt",
    (src:String,secret:String)
    	=>decrypt(src, secret),StringType)

三、在SparkSQL中调用函数

val frm: DataFrame = spark.createDataFrame(Seq(
	Test(1,Array("money","freedom"),Map("java"->85,"mysql"->67)),
  	Test(2,Array("beauty","beauty"),Map("java"->72,"mysql"->90)),
  	Test(3,Array("sports","beauty"),Map("java"->76,"html"->52))
))

val secret = "henryyb2211ariel"

val frmEncrypt: DataFrame = frm
  .select($"id",
	callUDF(
	  "aes_encrypt",
	  array_join($"hobbies", ","),
	  lit(secret)
	).as("encrypted_hobbies")
  )

val frmDecrypt: DataFrame = frmEncrypt
  .select($"id",
	split(
	  callUDF(
		"aes_decrypt",
		$"encrypted_hobbies",
		lit(secret)
	  ),
	  ","
	).as("hobbies")
  ).show()