# 創建 S3 存儲桶 
aws s3 mb s3://databrew-pii-data-zhanla --region cn-northwest-1

# 創建必要的文件夾（前綴） 
aws s3api put-object --bucket databrew-pii-data-zhanla --key sensitive_data_input/ --region cn-northwest-1
aws s3api put-object --bucket databrew-pii-data-zhanla --key profile_job_out/ --region cn-northwest-1
aws s3api put-object --bucket databrew-pii-data-zhanla --key encrypted_data_output/ --region cn-northwest-1

# 上傳樣例數據 
aws s3 cp ./sample-data.csv s3://databrew-pii-data-zhanla/sensitive_data_input/sample-data.csv --region cn-northwest-1

create database pii_sample_db

--Create table databrew_pii_sensitive_data from S3 file with csv format
CREATE EXTERNAL TABLE `databrew_pii_sensitive_data`(
`SSN` string, 
`gender` string, 
`birthdate` string, 
`maidenname` string, 
`lastname` string, 
`firstname` string, 
`address` string, 
`city` string, 
`state` string, 
`zip` string, 
`phone` string, 
`email` string, 
`cc_type` string, 
`CCN` string, 
`cc_cvc` string, 
`cc_expiredate` string
)
ROW FORMAT DELIMITED 
FIELDS TERMINATED BY ',' 
LOCATION
's3://databrew-pii-data-zhanla/sensitive_data_input/'
TBLPROPERTIES ("skip.header.line.count"="1");  

--Create table databrew_pii_encrypted_data from S3 file with Parquet format
CREATE EXTERNAL TABLE `databrew_pii_encrypted_data` (
`SSN` string, 
`gender` string, 
`birthdate` string, 
`maidenname` string, 
`lastname` string, 
`firstname` string, 
`address` string, 
`city` string, 
`state` string, 
`zip` int, 
`phone` string, 
`email` string, 
`cc_type` string, 
`CCN` string, 
`cc_cvc` int, 
`cc_expiredate` string
)
STORED AS PARQUET
LOCATION 's3://databrew-pii-data-zhanla/encrypted_data_output/'

package com.mycompany.tools.athena;

import java.util.*;
import java.util.Base64.*;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import org.apache.commons.lang3.StringUtils;
import com.amazonaws.athena.connector.lambda.handlers.UserDefinedFunctionHandler;
import com.amazonaws.encryptionsdk.AwsCrypto;
import com.amazonaws.encryptionsdk.CommitmentPolicy;
import com.amazonaws.encryptionsdk.CryptoResult;
import com.amazonaws.encryptionsdk.kmssdkv2.KmsMasterKey;
import com.amazonaws.encryptionsdk.kmssdkv2.KmsMasterKeyProvider;

/**
   * Athena PII 數據解密用户定義函數 (UDF)
   * 
   * 功能：在 Amazon Athena 中解密由 Amazon Glue DataBrew 加密的 PII 敏感數據字段
   * 用途：通過 Lambda 函數實現 Athena UDF，支持在 SQL 查詢中直接解密加密字段
   * 加密方式：使用 Amazon KMS 密鑰和 Amazon Encryption SDK 進行數據加密/解密
   */
  public class AthenaPIIUDF extends UserDefinedFunctionHandler
  {
      /** UDF 源類型標識 */
      private static final String SOURCE_TYPE = "MyCompany";
  
      /** Base64 解碼器，用於解碼加密數據 */
      private static Decoder base64Decoder = Base64.getDecoder();
  
      /** KMS 主密鑰提供者 */
      private static KmsMasterKeyProvider keyProvider = null;
  
    /** 靜態緩存映射，存儲不同 KMS 密鑰 ARN 對應的密鑰提供者，避免重複創建 */
      private static Map<String, KmsMasterKeyProvider> staticMap = new HashMap<String, KmsMasterKeyProvider>();

      /**加密 SDK 客户端，配置為禁止加密但允許解密 */
    private static final AwsCrypto crypto = AwsCrypto.builder()
            .withCommitmentPolicy(CommitmentPolicy.ForbidEncryptAllowDecrypt)
            .build();

      /**
       * 構造函數
       * 初始化 Athena UDF 處理器
       */
      public AthenaPIIUDF() {
          super(SOURCE_TYPE);
      }
  
      /**
       * 解密函數 - Athena UDF 主要功能
       * 
       * @param ciphertext 加密的密文（Base64 編碼）
       * @param keyArn KMS 密鑰 ARN，用於解密
       * @return 解密後的明文字符串
       * @throws IllegalStateException 當使用錯誤的密鑰時拋出異常
       */
      public static String decrypt(String ciphertext, String keyArn) {

        // 檢查輸入密文是否為空
        if (StringUtils.isBlank(ciphertext)) {
           return ciphertext;
        }

        // 移除字符串兩端的引號（如果存在）
        if(ciphertext.startsWith("\"")){
            ciphertext = ciphertext.substring(1,ciphertext.length() - 1);
        }

        // 從緩存中獲取或創建 KMS 密鑰提供者
        if(staticMap.containsKey(keyArn)){
            keyProvider = staticMap.get(keyArn);
        }else{
            // 創建新的 KMS 密鑰提供者並緩存
            keyProvider = KmsMasterKeyProvider.builder().buildStrict(keyArn);
            staticMap.put(keyArn, keyProvider);
        }

        // Base64 解碼密文
        byte[] byteContent = base64Decoder.decode(ciphertext);

        // 使用 Amazon Encryption SDK 解密數據
        final CryptoResult<byte[], KmsMasterKey> decryptResult = crypto.decryptData(keyProvider, byteContent);

        // 驗證解密使用的密鑰是否正確
        if (!decryptResult.getMasterKeyIds().get(0).equals(keyArn)) {
            throw new IllegalStateException("Wrong key id!");
        }

        // 將解密結果轉換為 UTF-8 字符串返回
            return new String(decryptResult.getResult(), StandardCharsets.UTF_8);
    }

USING EXTERNAL FUNCTION decrypt(t1 varchar, t2 varchar)
RETURNS varchar 
LAMBDA 'athena-udf-gluepii'
SELECT 
       decrypt(ssn, 'arn:aws-cn:kms:<region>:<account-id>:key/<key-id>')  ssnplaintext,*
FROM databrew_pii_encrypted_data