如何使用 AWS KMS 加密 Pandas/Spark 數據框中的列

Question

我想加密我的 Pandas（或 PySpark）數據幀的一列中的值，例如，獲取以下數據幀中的mobno列，對其進行加密並將結果放入encrypted_value列中：

我想使用 AWS KMS 加密密鑰。 我的問題是：如何實現這一目標的最優雅的方式是什么？

我正在考慮使用 UDF，它將調用 boto3 的 KMS 客戶端。 就像是：

@udf
def encrypt(plaintext):
  response = kms_client.encrypt(
    KeyId=aws_kms_key_id,
    Plaintext=plaintext
  )
  ciphertext = response['CiphertextBlob']
  return ciphertext

然后將此 udf 應用於整個列。

但我不太相信這是正確的方法。 這源於我是加密新手這一事實 - 首先，我什至不知道這個kms_client_encrypt函數是用於加密值（來自列）還是用於操作密鑰。 也許更好的方法是獲取密鑰，然后使用一些 python 加密庫（例如hashlib ）。

我想對加密過程進行一些澄清，並建議列加密的最佳方法是什么。

Answer 1

為避免在 UDF 中多次調用 KMS 服務，請改用 AWS Secrets Manager 來檢索您的加密密鑰並pycrypto來加密列。 以下作品：

from pyspark.sql.functions import udf, col
from Crypto.Cipher import AES

region_name = "eu-west-1"
session = boto3.session.Session()
client = session.client(service_name='secretsmanager', region_name=region_name)
get_secret_value_response = client.get_secret_value(SecretId=secret_name)
secret_key = json.loads(get_secret_value_response['SecretString'])
clear_text_column = 'mobo'

def encrypt(key, text):
    obj = AES.new(key, AES.MODE_CFB, 'This is an IV456')
    return obj.encrypt(text)

def udf_encrypt(key):
    return udf(lambda text: encrypt(key, text))

df.withColumn("encrypted", udf_encrypt(secret_key)(col(clear_text_column))).show()

或者，使用@Vektor88（PySpark 3 語法）建議的更高效的 Pandas UDF：

from functools import partial

encrypt_with_key = partial(encrypt, secret_key)

@pandas_udf(BinaryType())
def pandas_udf_encrypt(clear_strings: pd.Series) -> pd.Series:
    return clear_strings.apply(encrypt_with_key)

df.withColumn('encrypted', pandas_udf_encrypt(clear_text_column)).show()

Answer 2

從Spark 3.3開始，您可以在沒有 UDF 的情況下進行 AES 加密（和解密）。

aes_encrypt (expr, key[, mode[, padding]]) - 在給定mode使用 AES 和指定的padding返回expr的加密值。 支持 16、24 和 32 位的密鑰長度。 ( mode , padding ) 支持的組合是 ('ECB', 'PKCS') 和 ('GCM', 'NONE')。 默認模式為 GCM。

論據：
expr - 要加密的二進制值。
key - 用於加密數據的密碼。
mode - 指定應該使用哪種分組密碼模式來加密消息。 有效模式：ECB、GCM。
padding - 指定如何填充長度不是塊大小倍數的消息。 有效值：PKCS、NONE、DEFAULT。 DEFAULT 填充意味着 ECB 的 PKCS 和 GCM 的 NONE。

from pyspark.sql import functions as F
df = spark.createDataFrame([('8223344556',)], ['mobno'])

df = df.withColumn('encrypted_value', F.expr("aes_encrypt(mobno, 'your_secret_keyy')"))

df.show()
#  +----------+--------------------+
#  |     mobno|     encrypted_value|
#  +----------+--------------------+
#  |8223344556|[9B 33 DB 9B 5D C...|
#  +----------+--------------------+

df.printSchema()
#  root
#   |-- mobno: string (nullable = true)
#   |-- encrypted_value: binary (nullable = true)