如何使用 AWS KMS 加密 Pandas/Spark 数据帧中的列

Question

我想对我的 Pandas（或 py/spark）数据帧的一列中的值进行加密，例如，获取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()