DEV Community: Mostefa Brougui

Enhancing Data Security with Spark: A Guide to Column-Level Encryption - Part 2

Mostefa Brougui — Sun, 22 Dec 2024 17:19:50 +0000

This post describes how you can build an AWS Glue ingestion job with PySpark aes_encrypt() function to encrypt sensitive columns. It is part of a series that shows how column-level encryption can be deployed at scale using AWS Glue, AWS KMS and Amazon Athena or Amazon Redshift.

Introduction

In this post, I demonstrate setting up an AWS Glue ingestion job to encrypt sensitive columns using AWS KMS. We will also explore key management approaches and their impact on organizational security practices.

In the previous post, I introduced column-level encryption using a Jupyter Notebook and a static AWS KMS-generated data key. While useful for learning, hardcoding encryption keys is insecure and impractical for production. Instead, encryption keys should be dynamically managed and accessed securely at runtime.

AWS Glue is a "scalable, serverless data integration service that simplifies data discovery, preparation, and combination for analytics, machine learning, and application development." Glue's ability to attach IAM roles to jobs allows seamless interaction with other AWS services like S3 and KMS, enabling tasks such as data ingestion, manipulation (e.g., encryption/decryption), and storage.

We will also address key management questions: How should your Glue job encrypt sensitive columns? Where should encryption material be stored? Who should have access to it?

Let’s dive in!

Getting Started

Prepare Your Environment

To set up your AWS Glue for the first time, from the AWS Management Console:

Open the AWS Glue console and select Prepare your account for AWS Glue.
You can ignore Choose IAM users and roles for AWS Glue and choose Next if you'll be perfoming the next steps using your current role. Otherwise, select the IAM roles or users (I don't recommend using IAM users) that need to have access.
Under Grant Amazon S3 access, choose Next unless you want to edit the options selected by default. For the sake of this article, I granted access to all my S3 buckets. You will not do the same in a Production setting.
Under Choose a default service role, keep the default settings and choose Next, unless you have an existing IAM role for Glue.
Review and confirm your changes by choosing Apply changes.

To start building your ingestion job, create a Glue notebook by following the steps below.

From the AWS Glue console, on the left pane, under ETL jobs, choose Notebooks.
On the page Create job, select Notebook under Author using an interactive code notebook.
On the Notebook pop-up, keep the default settings and choose Create notebook.

AWS Glue will spin-up a Glue Studio notebook for you. Start by running the cell that initializes the Glue job and wait for Waiting for session <GUID> to get into ready status.... Your notebook is ready !

Encrypting Sensitive Columns with aes_encrypt()

Now that your Glue job is ready, let’s encrypt specific columns in your dataset using the PySpark function aes_encrypt(). We’ll build a reusable Python module, ColEncrypt, to handle both encryption and decryption, simplifying column-level encryption management in your Glue jobs.

Key Components

KeyManager: Manages the creation and decryption of Data Encryption Keys (DEKs) using AWS KMS.
ColEncrypt: Handles column encryption and decryption, leveraging PySpark's built-in AES functions.

These two components work together to provide a flexible, an almost production-ready solution for column-level encryption with better error handling and monitoring.

The source code is here.

Setting Up Key Management

Hardcoding encryption keys is a bad practice. Instead, use AWS KMS to generate and decrypt DEKs securely at runtime. Here’s how KeyManager handles this:

Here's how you can accomplish this in the class KeyManager.

class KeyManager:
    def __init__(self, kms_client):
        self.kms_client = kms_client

    def generate_data_key(self, key_id: str) -> bytes:
        try:
            response = self.kms_client.generate_data_key(KeyId=key_id, KeySpec="AES_256")
            return response["CiphertextBlob"]
        except ClientError as err:
            logger.error(f"Failed to generate data key: {err}")
            raise

    def decrypt_data_key(self, encrypted_key: bytes) -> bytes:
        try:
            response = self.kms_client.decrypt(CiphertextBlob=encrypted_key)
            return response["Plaintext"]
        except ClientError as err:
            logger.error(f"Failed to decrypt data key: {err}")
            raise

The generate_data_key() method fetches an encrypted DEK from KMS, while decrypt_data_key() decrypts it for use in encryption tasks. This ensures a secure, scalable, and auditable approach to key management.

Encrypting Columns with ColEncrypt

The ColEncrypt class applies AES encryption to specified columns, using DEKs managed by KeyManager.

`encrypt()` Method

This method handles the encryption process:

class ColEncrypt:
    def encrypt(self):
        key_manager = KeyManager(self.kms_client)
        for column in self.columns:
            dek = key_manager.generate_data_key(self.key_id)
            decrypted_dek = key_manager.decrypt_data_key(dek)
            dek_b64 = b64.b64encode(dek).decode("utf-8")

            self.df = self.df.withColumn(
                "key", lit(decrypted_dek)
            ).withColumn(
                column,
                concat(lit(dek_b64 + "::"), base64(expr(f"aes_encrypt({column}, key)")))
            ).drop("key")
        return self.df

Key Points

Base64 Encode Encrypted DEK:
- The encrypted DEK is encoded using Base64 for storage as a string.
- The prefix format is encrypted_dek_b64::encrypted_column_value.
Concatenation with Separator:
- The concat function adds the encrypted_dek_b64 followed by a separator (::) to the encrypted column value.
Temporary Key Column:
- A temporary key column is used to store the decrypted DEK during encryption, and it's dropped afterward.

Example Output

For a column named sensitive_column:

Original Value: 12345
Encrypted Column: ENCRYPTED_DEK_BASE64::ENCRYPTED_VALUE

Usage Example

Here’s how you can use ColEncrypt in your Glue job:

from ColEncrypt import ColEncrypt

# Load the data frame
df = glueContext.create_dynamic_frame.from_options(
    connection_type="s3",
    connection_options={"paths": ["s3://your-bucket/your-data.csv"]},
    format="csv",
    format_options={"withHeader": True}
).toDF()

# Specify columns to encrypt
kms_client = boto3.client("kms", region_name="us-east-2")
columns_to_encrypt = ["column_1",..., "column_n"]

# Initialize and encrypt
encryptor = ColEncrypt(df, columns_to_encrypt, "alias/your-kms-key", kms_client, "arn:aws:iam::123456789012:role/YourRole")
encrypted_df = encryptor.encrypt()

Initialize the ColEncrypt class with your DataFrame, sensitive columns, KMS Key ID, and KMS client or IAM Role if you want Glue to use different credentials to interact with the KMS key.
Call the encrypt() method with or without specifying additional columns.

This approach ensures the encrypted DEK is stored alongside the encrypted value, enabling efficient decryption during data processing. You can generate and/or fetch the DEKs from a centralized Key Store as well, which will not be covered in this post.

Decrypting Columns with ColEncrypt

To decrypt, extract the DEK from the first row of the encrypted column, decrypt it using KMS, and apply aes_decrypt().

`decrypt()` Method

class ColEncrypt:
    def decrypt(self):
        key_manager = KeyManager(self.kms_client)
        for column in self.columns:
            first_row = self.df.select(column).first()
            dek_b64, _ = first_row[column].split("::", 1)
            encrypted_dek = b64.b64decode(dek_b64)
            decrypted_dek = key_manager.decrypt_data_key(encrypted_dek)
            decrypted_dek_b64 = b64.b64encode(decrypted_dek).decode("utf-8")

            self.df = self.df.withColumn(
                column,
                expr(f"aes_decrypt(unbase64(split({column}, '::')[1]), unbase64('{decrypted_dek_b64}'))").cast("string")
            )
        return self.df

Usage Example

Here is an example of how you can decrypt a previously encrypted data frame using the same KMS key.

decryptor = ColEncrypt(encrypted_df, columns_to_encrypt, "alias/your-kms-key", kms_client)
decrypted_df = decryptor.decrypt()

With these methods, you can easily decrypt sensitive data columns for downstream processing or analysis.

Validating Data Integrity

After decryption, validate that the original and decrypted data match:

if original_df.subtract(decrypted_df).isEmpty() and decrypted_df.subtract(original_df).isEmpty():
    print("Original and decrypted files are identical!")
else:
    print("Original and decrypted files differ!")

This step ensures your encryption and decryption workflows are functioning correctly and that no data is lost or altered.

Conclusion

Column-level encryption with AWS Glue, PySpark, and AWS KMS ensures secure data handling. By implementing ColEncrypt, you can streamline encryption and decryption while adhering to best practices for key management. This approach not only secures sensitive data but also provides a scalable, auditable framework for enterprise-grade security.

Stay tuned for the last part of this series, where I'll explore secure analytics on encrypted data with Amazon Redshift.

Enhancing Data Security with Spark: A Guide to Column-Level Encryption - Part 1

Mostefa Brougui — Fri, 08 Mar 2024 16:26:29 +0000

This post describes how you can use PySpark aes_encrypt() function to encrypt sensitive columns when ingesting data. It is part of a series that shows how column-level encryption can be deployed at scale using AWS Glue, AWS KMS and Amazon Athena or Amazon Redshift.

Introduction

In an era where data breaches are increasingly common, securing sensitive data is not just a best practice but a necessity. As Werner Vogels, Amazon's CTO, wisely put it: "Dance Like Nobody’s Watching. Encrypt Like Everyone Is." In this series, we explore how to deploy column-level encryption at scale using an integration of PySpark, AWS Glue, AWS KMS, and Amazon Redshift or Amazon Athena.

Why Column-Level Encryption?

Column-level encryption allows for fine-grained control over data access. It's particularly useful in scenarios where different groups require access to specific data subsets. This post serves as an introductory guide, focusing on the implementation of PySpark's aes_encrypt() function for encrypting sensitive data during ingestion.

Setting the Stage

Before diving into the encryption process, it's important to understand the relevant technologies:

PySpark: A Python API for Apache Spark, used for large-scale data processing.
AWS Glue: A serverless data integration service that makes it easy to discover, prepare, and combine data for analytics, machine learning, and application development.
AWS KMS: AWS Key Management Service, a managed service that makes it easy to create and manage cryptographic keys.
Amazon Athena: An interactive query service that makes it easy to analyze data in Amazon S3 using standard SQL.
Amazon Redshift: A fast, fully managed, petabyte-scale data warehouse service that makes it simple and cost-effective to efficiently analyze all your data.

Getting Started

Prepare Your Environment

Run a Jupyter Notebook on Docker. Here's how.
Clone this repository for sample data and scripts: GitHub Link.

Read Sample Data

Generate / Upload a CSV file containing a data sample to work with to the notebook. I provide a sample file in the provided GitHub repository: CSV file.
Read the sample data using the code below.

import base64 as b64
from pyspark.sql import SparkSession
from pyspark.sql.types import *

spark = SparkSession.builder.appName("MyApp").getOrCreate()

df = spark.read.option("header", True).csv("sample-pci-data.csv")
df.show()

You should see something similar to the below.

+-------------------+-----------+-------------------+
|First and Last Name|        SSN| Credit Card Number|
+-------------------+-----------+-------------------+
|      Robert Aragon|489-36-8350|4929-3813-3266-4295|
|      Ashley Borden|514-14-8905|5370-4638-8881-3020|
...
only showing top 2 rows

Understanding aes_encrypt()

The aes_encrypt() function in PySpark 3.3 offers AES encryption with various configurations. Read more about it here.

Below is the syntax of the function. It uses AES encryption with 16, 24, or 32-bit keys in ECB, GCM, or CBC modes with appropriate padding. IVs are optional for CBC and GCM modes and are auto-generated if absent. For GCM, optional AAD must match during encryption and decryption. GCM is the default mode.

def aes_encrypt(
    input: "ColumnOrName",
    key: "ColumnOrName",
    mode: Optional["ColumnOrName"] = None,
    padding: Optional["ColumnOrName"] = None,
    iv: Optional["ColumnOrName"] = None,
    aad: Optional["ColumnOrName"] = None,
) -> Column:

The following exanple shows the encryption and decryption of the string Spark SQL using the key abcdefghijklmnop which is a 16 characters long = 16 bytes = 128 bits key.

df = spark.createDataFrame([(
    "Spark SQL", "abcdefghijklmnop",)],
    ["input", "key"]
)

df.select(aes_decrypt(
    unbase64(base64(aes_encrypt(df.input, df.key))), df.key
).cast("STRING").alias('r')).collect()

[Row(r='Spark SQL')]

Generating an Encryption Key

Use AWS KMS's GenerateDataKey API for a robust key generation. This requires an existing KMS key.
- To generate the key using the AWS CLI generate-data-key command, use the example shown below. Make sure to replace AWS_PROFILE and KMS_KEY_ID with the correct values.
```
aws kms generate-data-key --profile AWS_PROFILE --key-id KMS_KEY_ID --key-spec AES_256
```
The response shown below contains the key in Plaintext base64-encoded format and in encrypted (CiphertextBlob) format.
```
{
    "CiphertextBlob": "AQIDAHjhvhi8z21Psjp5qjRibLXgHGkMMP/BmNWdIKTuZnLdNwHHYQjpmQo0oyk5rk07mjZTAAAAfjB8BgkqhkiG9w0BBwagbzBtAgEAMGgGCSqGSIb3DQEHATAeBglghkgBZQMEAS4wEQQM6uOfZ7iyJ2UztfidAgEQgDuFQ5fajsvnhfaCmN/q8kk1JinY7gHqT/Bz9W3RxqkEv5ggPZKELcQtqNbRYdEzSwKPDTs+Grp0RtRqWw==",
    "Plaintext": "tEZPizBEj5EG5IDY1SvAECa5yZa5fVP1SrJGsGimx9I=",
    "KeyId": "KMS_KEY_ID"
}
```
- If you generate the data key using boto3 kms client, you will get a byte string, similar to b"\xc27~\x15\xc7\x9a\x8a|\xb48\\\xd7\x894g-v\xac\xb5\n%\x17\x96g\xab\x88\x8a;|bU/". This value can be used as-is with the aes_encrypt() function.
Alternatively, derive a key from a string using PBKDF2. (Link to code sample)

Encrypting Sensitive Columns

We'll encrypt the "SSN" and "Credit Card Number" columns as shown below.

encryption_key = "tEZPizBEj5EG5IDY1SvAECa5yZa5fVP1SrJGsGimx9I="
# Decode the key from Base64 
decoded_key = b64.b64decode(encryption_key)
# Encrypt two columns at the "same time" using the same key
df_encrypted = df.withColumn('SSN_Encrypted', base64(expr(f"aes_encrypt(SSN, unhex('{decoded_key.hex()}'), 'GCM')")))\ 
                 .withColumn('CreditCardNumber_Encrypted', base64(expr(f"aes_encrypt(Credit_Card_Number, unhex('{decoded_key.hex()}'), 'GCM')")))
df_encrypted.show(truncate=True)

# +----------------+-----------+-------------------+--------------------+--------------------------+
# | First_Last_Name|        SSN| Credit_Card_Number|       SSN_Encrypted|CreditCardNumber_Encrypted|
# +----------------+-----------+-------------------+--------------------+--------------------------+
# |   Robert Aragon|489-36-8350|4929-3813-3266-4295|HI39DkVYN9WLTWAH3...|      XqvsWgRty1CBkJ7c9...|
# |   Ashley Borden|514-14-8905|5370-4638-8881-3020|5ME9bErsff7Zhzw5z...|      mhizKvf5053KqStxp...|
#  ...
# only showing top 2 rows

# Output path to store the new CSV
output_path = "./encrypted/sample-pci-data-encrypted.csv"
df_encrypted.write.csv(path=output_path, mode="overwrite", header=True, sep=",")

Decryption Process

Here's how you can decrypt the data using the same key.

df = spark.read.option("header", True).csv("encrypted/sample-pci-data-encrypted.csv")
df.show()
df_decrypted = df.withColumn('SSN_Decrypted', expr(f"aes_decrypt(unbase64(SSN_Encrypted), unhex('{decoded_key.hex()}'), 'GCM')").cast("STRING"))
                 .withColumn('CC_Decrypted', expr(f"aes_decrypt(unbase64(CreditCardNumber_Encrypted), unhex('{decoded_key.hex()}'), 'GCM')").cast("STRING"))

# +----------------+-----------+-------------------+--------------------+--------------------------+
# | First_Last_Name|        SSN| Credit_Card_Number|       SSN_Encrypted|CreditCardNumber_Encrypted|
# +----------------+-----------+-------------------+--------------------+--------------------------+
# |   Robert Aragon|489-36-8350|4929-3813-3266-4295|Ft4y7DmP8xN8QtdHn...|      4f7Usr2dvVWbG26Bd...|
# |   Ashley Borden|514-14-8905|5370-4638-8881-3020|nsdeq//on5T/L64Ow...|      8LN7pIP1pa7Nmoma1...|
# only showing top 2 rows

# +----------------+-----------+-------------------+--------------------+--------------------------+-------------+-------------------+
# | First_Last_Name|        SSN| Credit_Card_Number|       SSN_Encrypted|CreditCardNumber_Encrypted|SSN_Decrypted|       CC_Decrypted|
# +----------------+-----------+-------------------+--------------------+--------------------------+-------------+-------------------+
# |   Robert Aragon|489-36-8350|4929-3813-3266-4295|Ft4y7DmP8xN8QtdHn...|      4f7Usr2dvVWbG26Bd...|  489-36-8350|4929-3813-3266-4295|
# |   Ashley Borden|514-14-8905|5370-4638-8881-3020|nsdeq//on5T/L64Ow...|      8LN7pIP1pa7Nmoma1...|  514-14-8905|5370-4638-8881-3020|
# only showing top 2 rows

Conclusion

In this post, I walked through the steps of encrypting and decrypting sensitive data columns using PySpark in a Jupyter Notebook. This setup serves as a foundational step towards a more complex and scalable data ingestion and consumption model, which I'll explore in upcoming posts, including the integration of AWS Glue and AWS KMS.

Stay tuned for further insights on scaling this approach and managing encryption keys with AWS KMS!

DEV Community: Mostefa Brougui

Enhancing Data Security with Spark: A Guide to Column-Level Encryption - Part 2

Introduction

Getting Started

Prepare Your Environment

Encrypting Sensitive Columns with aes_encrypt()

Key Components

Setting Up Key Management

Encrypting Columns with ColEncrypt

encrypt() Method

Key Points

Example Output

Usage Example

Decrypting Columns with ColEncrypt

decrypt() Method

Usage Example

Validating Data Integrity

Conclusion

Enhancing Data Security with Spark: A Guide to Column-Level Encryption - Part 1

Introduction

Why Column-Level Encryption?

Setting the Stage

Getting Started

Prepare Your Environment

Read Sample Data

Understanding aes_encrypt()

Generating an Encryption Key

Encrypting Sensitive Columns

Decryption Process

Conclusion

`encrypt()` Method

`decrypt()` Method