Use KMS to build MaxCompute UDFs to encrypt and decrypt data

<dependencies>
        <dependency>
            <groupId>com.aliyun.odps</groupId>
            <artifactId>odps-sdk-udf</artifactId>
            <version>0.48.6-public</version>
        </dependency>
        <dependency>
            <groupId>com.aliyun</groupId>
            <artifactId>alibabacloud-kms-java-sdk</artifactId>
            <version>1.2.5</version>
            <exclusions>
                <exclusion>
                    <groupId>com.aliyun</groupId>
                    <artifactId>tea</artifactId>
                </exclusion>
            </exclusions>
        </dependency>
        <dependency>
            <groupId>com.aliyun</groupId>
            <artifactId>tea</artifactId>
            <version>1.2.3</version>
        </dependency>
</dependencies>

package com.aliyun.kms.sample;

import com.aliyun.tea.utils.StringUtils;

import java.nio.charset.StandardCharsets;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class UDFUtils {

    private UDFUtils() {
        // do noting
    }

    public static boolean isNull(String value) {
        return StringUtils.isEmpty(value) || "\\N".equals(value);
    }

    public static String getSynchronousKeyVersion(String udfIdentify, String keyIdentify) {
        UUID uuid = null;
        byte[] bytes = null;
        if (StringUtils.isEmpty(keyIdentify)) {
            bytes = udfIdentify.getBytes(StandardCharsets.UTF_8);
            uuid = UUID.nameUUIDFromBytes(bytes);
            return uuid.toString();
        }
        bytes = (udfIdentify + "/" + keyIdentify).getBytes(StandardCharsets.UTF_8);
        uuid = UUID.nameUUIDFromBytes(bytes);
        return uuid.toString();
    }

    public static <K, V> Map<K, V> getLRUMap(int capacity) {
        return new LRUMap<K, V>(capacity);
    }

    private static class LRUMap<K, V> extends LinkedHashMap<K, V> {
        private final int capacity;
        private final ReadWriteLock readWriteLock = new ReentrantReadWriteLock();
        private final Lock readLock = readWriteLock.readLock();
        private final Lock writeLock = readWriteLock.writeLock();

        public LRUMap(int capacity) {
            super(capacity, 0.75f, true);
            this.capacity = capacity;
        }

        @Override
        public V get(Object key) {
            readLock.lock();
            try {
                return super.get(key);
            } finally {
                readLock.unlock();
            }
        }

        @Override
        public V put(K key, V value) {
            writeLock.lock();
            try {
                return super.put(key, value);
            } finally {
                writeLock.unlock();
            }
        }

        @Override
        protected boolean removeEldestEntry(Map.Entry<K, V> eldest) {
            return size() > capacity;
        }
    }

}

package com.aliyun.kms.sample;

public class UDFConstant {

    private UDFConstant() {
        // do nothing
    }
    public static final String KMS_UDF_CONF_NAME = "kms_udf_conf.properties";
    public static final String UDF_ENCRYPT_KEY_ID_KEY = "udf.kms.keyId";
    public static final String KMS_CLIENT_KEY_FILE_KEY = "udf.kms.clientkey.file";
    public static final String KMS_CLIENT_KEY_PASSWORD_KEY = "udf.kms.clientkey.password";
    public static final String KMS_INSTANCE_CA_FILE_KEY = "udf.kms.instance.ca.file";
    public static final String KMS_INSTANCE_ENDPOINT_KEY = "udf.kms.instance.endpoint";
    public static final String KMS_SYNCHRONOUS_SECRET_NAME_KEY = "udf.synchronous.secret.name";
    public static final String KMS_RAM_SECRET_NAME_KEY = "udf.ram.secret.name";
    public static final String KMS_ENDPOINT_KEY = "udf.kms.endpoint";
    public static final String ENCRYPT_DATA_KEY_FORMAT = "%04d";
    public static final int GCM_IV_LENGTH = 12;
    public static final int GCM_IV_BASE64_LENGTH = 16;
    public static final int GCM_TAG_LENGTH = 16;
}

package com.aliyun.kms.sample;

import com.aliyun.dkms.gcs.openapi.models.Config;
import com.aliyun.kms.kms20160120.TransferClient;
import com.aliyun.kms.kms20160120.model.KmsConfig;
import com.aliyun.kms.kms20160120.model.KmsRuntimeOptions;
import com.aliyun.kms20160120.Client;
import com.aliyun.kms20160120.models.DecryptRequest;
import com.aliyun.kms20160120.models.DecryptResponse;
import com.aliyun.odps.udf.ExecutionContext;
import com.aliyun.odps.udf.UDF;
import com.aliyun.odps.udf.UDFException;
import org.apache.commons.codec.binary.Base64;

import javax.crypto.Cipher;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.Properties;

public class DecryptUDF extends UDF {

    private static Client client;
    private static Map<String, byte[]> dataKeyMap;
    private static final int BASE_STEP=4;

    private static Base64 base64 = new Base64();

    @Override
    public void setup(ExecutionContext ctx) throws UDFException, IOException {
        super.setup(ctx);
        try {
            InputStream in = ctx.readResourceFileAsStream(UDFConstant.KMS_UDF_CONF_NAME);
            Properties properties = new Properties();
            properties.load(in);
            KmsConfig config = new KmsConfig();
            String clientKeyFileName = properties.getProperty(UDFConstant.KMS_CLIENT_KEY_FILE_KEY);
            byte[] clientKeyFileContent = ctx.readResourceFile(clientKeyFileName);
            config.setClientKeyContent(new String(clientKeyFileContent, StandardCharsets.UTF_8));
            String caFileName = properties.getProperty(UDFConstant.KMS_INSTANCE_CA_FILE_KEY);
            byte[] caFileContent = ctx.readResourceFile(caFileName);
            config.setCa(new String(caFileContent,StandardCharsets.UTF_8));
            config.setPassword(properties.getProperty(UDFConstant.KMS_CLIENT_KEY_PASSWORD_KEY));
            config.setEndpoint(properties.getProperty(UDFConstant.KMS_INSTANCE_ENDPOINT_KEY));
            dataKeyMap = UDFUtils.getLRUMap(1000);
            client = new TransferClient(config);
        } catch (Exception e) {
            throw new UDFException(e);
        }
    }

    public String evaluate(String data) {
        if (UDFUtils.isNull(data)) {
            return data;
        }
        if (data.length() <= BASE_STEP) {
            return data;
        }
        int dataLength = Integer.valueOf(data.substring(0, 4));
        if (data.length() <= BASE_STEP + dataLength) {
            return data;
        }
        String dataKeyCiphertext = data.substring(4, 4 + dataLength);
        try {

            byte[] dataKeyBytes = getDataKeyPlaintext(dataKeyCiphertext);
            Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
            SecretKeySpec keySpec = new SecretKeySpec(dataKeyBytes, "AES");
            byte[] iv = base64.decode(data.substring(4 + dataLength, 4 + dataLength + UDFConstant.GCM_IV_BASE64_LENGTH));
            byte[] cipherText = base64.decode(data.substring(4 + dataLength + UDFConstant.GCM_IV_BASE64_LENGTH));
            GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(UDFConstant.GCM_TAG_LENGTH * 8, iv);
            cipher.init(Cipher.DECRYPT_MODE, keySpec, gcmParameterSpec);
            byte[] decryptedData = cipher.doFinal(cipherText);
            return new String(decryptedData, "UTF-8");

        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

    private byte[] getDataKeyPlaintext(String dataKeyCiphertext) throws Exception {
        if (!dataKeyMap.containsKey(dataKeyCiphertext)) {
            synchronized (dataKeyCiphertext.intern()) {
                DecryptRequest request = new DecryptRequest();
                request.ciphertextBlob = dataKeyCiphertext;
                KmsRuntimeOptions runtimeOptions = new KmsRuntimeOptions();
                DecryptResponse response = client.decryptWithOptions(request, runtimeOptions);
                if (!dataKeyMap.containsKey(dataKeyCiphertext)) {
                    dataKeyMap.put(dataKeyCiphertext, base64.decode(response.getBody().plaintext));
                }
            }
        }
        return dataKeyMap.get(dataKeyCiphertext);
    }


    @Override
    public void close() throws UDFException, IOException {
        super.close();
    }
}

# Add the JAR package of the UDF implementation class to the MaxCompute project.
ADD JAR encrypt_udf.jar;
ADD JAR sync_key_encrypt_udf.jar;
ADD JAR decrypt_udf.jar;
# Add the configuration file to the MaxCompute project.
ADD FILE kms_udf_conf.properties;
# Add the client key file of the KMS instance to the MaxCompute project.
ADD FILE clientkey_xxxx.json;
# Add the CA certificate file of the KMS instance to the MaxCompute project.
ADD FILE PrivateKmsCA_xxxx.pem;
# Create an encryption UDF named encrypt_udf.
CREATE FUNCTION encrypt_udf AS 'com.aliyun.kms.sample.EncryptUDF'
USING 'encrypt_udf.jar,kms_udf_conf.properties,clientkey_xxxx.json,PrivateKmsCA_xxxx.pem';
# Create a high-performance encryption UDF named sync_key_encrypt_udf.
CREATE FUNCTION sync_key_encrypt_udf AS 'com.aliyun.kms.sample.SynchronousKeyEncryptUDF'
USING 'sync_key_encrypt_udf.jar,kms_udf_conf.properties,clientkey_xxxx.json,PrivateKmsCA_xxxx.pem';
# Create a decryption UDF named decrypt_udf. 
CREATE FUNCTION decrypt_udf AS 'com.aliyun.kms.sample.DecryptUDF'
USING 'decrypt_udf.jar,kms_udf_conf.properties,clientkey_xxxx.json,PrivateKmsCA_xxxx.pem';

-- Use a UDF to encrypt data.
SELECT encrypt_udf(column_name) FROM my_table;
-- Use a UDF to encrypt data. The value of version is the secret version, and the default value of version is null.
SELECT sync_key_encrypt_udf(version,column_name) FROM my_table;
-- Use a UDF to decrypt data.
SELECT decrypt_udf(column_name) FROM my_table;

pyodps==0.11.6.2
alibabacloud-kms-python-sdk==1.1.3

# coding=utf-8
import uuid
from collections import OrderedDict


def is_null(value):
    return len(value) == 0 or value == '\\N'


def generate_uuid(key_version):
    return uuid.uuid5(uuid.NAMESPACE_DNS, key_version)


def get_config_value(config_dict, dict_key):
    if dict_key not in config_dict:
        raise ValueError("can not find key[{}]" % dict_key)
    return config_dict[dict_key]


def properties_to_dict(cache_file):
    config_dict = {}
    for line in cache_file:
        line = line.strip()
        if not line:
            continue
        k, v = line.split("=")
        config_dict[k.strip()] = v.strip()
    return config_dict


def object_to_dict(obj):
    return obj.__dict__


class LRUCache(object):
    def __init__(self, capacity):
        self.capacity = capacity
        self.cache = OrderedDict()

    def get(self, key):
        if key not in self.cache:
            return -1
        value = self.cache.pop(key)  # Remove the item to update the order
        self.cache[key] = value  # Re-insert it to the end
        return value

    def put(self, key, value):
        if key in self.cache:
            self.cache.pop(key)  # Remove the existing entry to update the order
        elif len(self.cache) == self.capacity:
            self.cache.popitem(last=False)  # Remove the least recently used item
        self.cache[key] = value

    def contains_key(self, key):
        return key in self.cache

# coding=utf-8
UDF_CONFIG_CACHE_FILE = "kms_udf_conf.properties"
KMS_CLIENT_KEY_FILE_KEY = "udf.kms.clientkey.file"
KMS_CLIENT_KEY_PASSWORD_KEY = "udf.kms.clientkey.password"
KMS_INSTANCE_ENDPOINT_KEY = "udf.kms.instance.endpoint"
KMS_INSTANCE_CA_FILE_KEY = "udf.kms.instance.ca.file"
KMS_SYNCHRONOUS_SECRET_NAME_KEY = "udf.synchronous.secret.name"
KMS_RAM_SECRET_NAME_KEY = "udf.ram.secret.name"
KMS_REGION_KEY = "udf.kms.regionId"
UDF_ENCRYPT_KEY_ID_KEY = "udf.kms.keyId"
ENCRYPT_DATA_KEY_FORMAT = "%04d"
GCM_IV_LENGTH = 12
GCM_TAG_LENGTH = 16
DEFAULT_NUMBER_OF_BYTES = 32

# coding=utf-8
import base64

from alibabacloud_kms20160120.models import GenerateDataKeyRequest
from alibabacloud_kms_kms20160120.models import KmsConfig, KmsRuntimeOptions
from alibabacloud_kms_kms20160120.transfer_client import TransferClient
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from odps.distcache import get_cache_file
import random
import string
import kms_udf_constant
import kms_udf_utils


class EncryptUDF(object):

    def __init__(self):
        cache_file = get_cache_file(kms_udf_constant.UDF_CONFIG_CACHE_FILE)
        self.config_dict = kms_udf_utils.properties_to_dict(cache_file)
        cache_file.close()
        kms_config = KmsConfig(
            protocol='https',
            client_key_file=kms_udf_utils.get_config_value(self.config_dict, kms_udf_constant.KMS_CLIENT_KEY_FILE_KEY),
            password=kms_udf_utils.get_config_value(self.config_dict, kms_udf_constant.KMS_CLIENT_KEY_PASSWORD_KEY),
            endpoint=kms_udf_utils.get_config_value(self.config_dict, kms_udf_constant.KMS_INSTANCE_ENDPOINT_KEY),
        )
        self.ca_file_path = kms_udf_utils.get_config_value(self.config_dict, kms_udf_constant.KMS_INSTANCE_CA_FILE_KEY)
        self.keyId = kms_udf_utils.get_config_value(self.config_dict, kms_udf_constant.UDF_ENCRYPT_KEY_ID_KEY)
        self.instance_client = TransferClient(kms_config=kms_config)
        resp = self.generate_data_key()
        self.encrypted_data_key_part = (kms_udf_constant.ENCRYPT_DATA_KEY_FORMAT %
                                        len(resp.body.ciphertext_blob) + resp.body.ciphertext_blob)
        self.plaintext = resp.body.plaintext
        self.ciphertext_blob = resp.body.ciphertext_blob

    def evaluate(self, data):
        if kms_udf_utils.is_null(data):
            return data
        iv = bytes(''.join(random.sample(string.ascii_letters + string.digits, kms_udf_constant.GCM_IV_LENGTH)),
                   encoding="utf-8")
        data_bytes = data.encode("utf-8")
        encryptor = Cipher(
            algorithms.AES(base64.b64decode(self.plaintext)),
            modes.GCM(iv),
        ).encryptor()
        ciphertext = encryptor.update(data_bytes) + encryptor.finalize()
        return self.encrypted_data_key_part + str(base64.b64encode(iv), "utf-8") + str(
            base64.b64encode(ciphertext + encryptor.tag), "utf-8")

    def generate_data_key(self):
        request = GenerateDataKeyRequest()
        request.key_id = self.keyId
        request.number_of_bytes = kms_udf_constant.DEFAULT_NUMBER_OF_BYTES
        runtime = KmsRuntimeOptions(
            ca=self.ca_file_path
        )
        return self.instance_client.generate_data_key_with_options(request, runtime)

# coding=utf-8
import base64

from alibabacloud_kms20160120 import models as kms_20160120_models
from alibabacloud_kms_kms20160120.models import KmsConfig, KmsRuntimeOptions
from alibabacloud_kms_kms20160120.transfer_client import TransferClient
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from odps.distcache import get_cache_file

import kms_udf_constant
import kms_udf_utils
from kms_udf_utils import LRUCache


class DecryptUDF(object):

    def __init__(self):
        cache_file = get_cache_file(kms_udf_constant.UDF_CONFIG_CACHE_FILE)
        self.config_dict = kms_udf_utils.properties_to_dict(cache_file)
        cache_file.close()
        kms_config = KmsConfig(
            protocol='https',
            client_key_file=kms_udf_utils.get_config_value(self.config_dict, kms_udf_constant.KMS_CLIENT_KEY_FILE_KEY),
            password=kms_udf_utils.get_config_value(self.config_dict, kms_udf_constant.KMS_CLIENT_KEY_PASSWORD_KEY),
            endpoint=kms_udf_utils.get_config_value(self.config_dict, kms_udf_constant.KMS_INSTANCE_ENDPOINT_KEY),
        )
        self.ca_file_path = kms_udf_utils.get_config_value(self.config_dict, kms_udf_constant.KMS_INSTANCE_CA_FILE_KEY)
        self.keyId = kms_udf_utils.get_config_value(self.config_dict, kms_udf_constant.UDF_ENCRYPT_KEY_ID_KEY)
        self.instance_client = TransferClient(kms_config=kms_config)
        self.data_key_cache = LRUCache(1000)

    def evaluate(self, data):
        if kms_udf_utils.is_null(data):
            return data
        if len(data) < 4:
            return data
        data_key_len = int(data[0:4])
        if len(data) < 4 + data_key_len:
            return data
        data_key_ciphertext = data[4:4 + data_key_len]
        if not self.data_key_cache.contains_key(data_key_ciphertext):
            decrypt_request = kms_20160120_models.DecryptRequest(
                ciphertext_blob=data_key_ciphertext
            )
            decrypt_runtime = KmsRuntimeOptions(
                ca=self.ca_file_path
            )
            decrypt_response = self.instance_client.decrypt_with_options(decrypt_request, decrypt_runtime)
            self.data_key_cache.put(data_key_ciphertext, base64.b64decode(decrypt_response.body.plaintext))
        data_key = self.data_key_cache.get(data_key_ciphertext)
        envelope_cipher_bytes = base64.b64decode(data[4 + data_key_len:])
        iv = envelope_cipher_bytes[0:kms_udf_constant.GCM_IV_LENGTH]
        cipher_bytes = envelope_cipher_bytes[
                   kms_udf_constant.GCM_IV_LENGTH:len(envelope_cipher_bytes) - kms_udf_constant.GCM_TAG_LENGTH]
        tag = envelope_cipher_bytes[len(envelope_cipher_bytes) - kms_udf_constant.GCM_TAG_LENGTH:]
        decryptor = Cipher(algorithms.AES(data_key), modes.GCM(iv, tag)).decryptor()
        return str(decryptor.update(cipher_bytes) + decryptor.finalize())

# Add the preceding program and requirements.txt-related dependencies to the resource.
ADD PY decrypt_udf.py;
ADD PY encrypt_udf.py;
# Add the configuration file to the MaxCompute project.
ADD FILE kms_udf_conf.properties;
# Add the client key file of the KMS instance to the MaxCompute project.
ADD FILE clientkey_xxxx.json;
# Add the CA certificate file of the KMS instance to the MaxCompute project.
ADD FILE PrivateKmsCA_xxxx.pem;
# Create an encryption UDF.
CREATE FUNCTION encrypt_udf AS 'com.aliyun.kms.sample.EncryptUDF'
USING 'encrypt_udf.py,kms_udf_conf.properties,clientkey_xxxx.json,PrivateKmsCA_xxxx.pem';
# Create a decryption UDF.
CREATE FUNCTION decrypt_udf AS 'com.aliyun.kms.sample.DecryptUDF'
USING 'decrypt_udf.py,kms_udf_conf.properties,clientkey_xxxx.json,PrivateKmsCA_xxxx.pem';

# Configure variables.
# Specify the key ID.
udf.kms.keyId=key-xxxxxxxxxxxx
# Specify the name of the client key file.
udf.kms.clientkey.file=clientkey_xxxx.json
# Specify the password of the client key.
udf.kms.clientkey.password=xxxxxxxx
# Specify the CA certificate of the KMS instance.
udf.kms.instance.ca.file=PrivateKmsCA_xxxx.pem
# Specify the endpoint of the KMS instance.
udf.kms.instance.endpoint=xxxxx.cryptoservice.kms.aliyuncs.com

-- Use a UDF to encrypt data.
SELECT encrypt_udf(column_name) FROM my_table;
-- Use a UDF to decrypt data.
SELECT decrypt_udf(column_name) FROM my_table;