# Copyright 2024 Confluent Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import base64 import io import logging import time from typing import Optional, Tuple, Any, List from tink import aead, daead, KmsClient, kms_client_from_uri, \ register_kms_client, TinkError from tink.core import Registry from tink.proto import tink_pb2, aes_siv_pb2 from confluent_kafka.schema_registry import SchemaRegistryError, RuleMode, \ _MAGIC_BYTE_V0 from confluent_kafka.schema_registry.rule_registry import RuleRegistry from confluent_kafka.schema_registry.rules.encryption.dek_registry.dek_registry_client import \ DekRegistryClient, Kek, KekId, DekId, Dek, DekAlgorithm from confluent_kafka.schema_registry.rules.encryption.kms_driver_registry import \ get_kms_driver, KmsDriver from confluent_kafka.schema_registry.serde import RuleContext, \ RuleError, RuleExecutor, FieldType, FieldRuleExecutor, FieldTransform, \ FieldContext log = logging.getLogger(__name__) aead.register() daead.register() ENCRYPT_KEK_NAME = "encrypt.kek.name" ENCRYPT_KMS_KEY_ID = "encrypt.kms.key.id" ENCRYPT_KMS_TYPE = "encrypt.kms.type" ENCRYPT_DEK_ALGORITHM = "encrypt.dek.algorithm" ENCRYPT_DEK_EXPIRY_DAYS = "encrypt.dek.expiry.days" ENCRYPT_ALTERNATE_KMS_KEY_IDS = "encrypt.alternate.kms.key.ids" MILLIS_IN_DAY = 24 * 60 * 60 * 1000 class Clock(object): def now(self) -> int: return int(round(time.time() * 1000)) class EncryptionExecutor(RuleExecutor): def __init__(self, clock: Clock = Clock()): self.client = None self.config = None self.clock = clock def configure(self, client_conf: dict, rule_conf: dict): if client_conf: if self.client: if self.client.config() != client_conf: raise RuleError("executor already configured") else: self.client = DekRegistryClient.new_client(client_conf) if self.config: if rule_conf: for key, value in rule_conf.items(): v = self.config.get(key) if v is not None: if v != value: raise RuleError(f"rule config key already set: {key}") else: self.config[key] = value else: self.config = rule_conf if rule_conf else {} def type(self) -> str: return "ENCRYPT_PAYLOAD" def transform(self, ctx: RuleContext, message: Any) -> Any: executor = self.new_transform(ctx) return executor.transform(ctx, FieldType.BYTES, message) def new_transform(self, ctx: RuleContext) -> 'EncryptionExecutorTransform': cryptor = self._get_cryptor(ctx) kek_name = self._get_kek_name(ctx) dek_expiry_days = self._get_dek_expiry_days(ctx) transform = EncryptionExecutorTransform( self, cryptor, kek_name, dek_expiry_days) return transform def close(self): if self.client is not None: self.client.__exit__() def _get_cryptor(self, ctx: RuleContext) -> 'Cryptor': dek_algorithm = DekAlgorithm.AES256_GCM dek_algorithm_str = ctx.get_parameter(ENCRYPT_DEK_ALGORITHM) if dek_algorithm_str is not None: dek_algorithm = DekAlgorithm[dek_algorithm_str] cryptor = Cryptor(dek_algorithm) return cryptor def _get_kek_name(self, ctx: RuleContext) -> str: kek_name = ctx.get_parameter(ENCRYPT_KEK_NAME) if kek_name is None: raise RuleError("no kek name found") if kek_name == "": raise RuleError("empty kek name") return kek_name def _get_dek_expiry_days(self, ctx: RuleContext) -> int: dek_expiry_days_str = ctx.get_parameter(ENCRYPT_DEK_EXPIRY_DAYS) if dek_expiry_days_str is None: return 0 try: dek_expiry_days = int(dek_expiry_days_str) except ValueError: raise RuleError("invalid expiry days") if dek_expiry_days < 0: raise RuleError("negative expiry days") return dek_expiry_days @classmethod def register(cls): RuleRegistry.register_rule_executor(EncryptionExecutor()) @classmethod def register_with_clock(cls, clock: Clock) -> 'EncryptionExecutor': executor = EncryptionExecutor(clock) RuleRegistry.register_rule_executor(executor) return executor class Cryptor: EMPTY_AAD = b"" def __init__(self, dek_format: DekAlgorithm): self.dek_format = dek_format self.is_deterministic = dek_format == DekAlgorithm.AES256_SIV self.registry = Registry() if dek_format is DekAlgorithm.AES128_GCM: self.key_template = aead.aead_key_templates.AES128_GCM_RAW elif dek_format is DekAlgorithm.AES256_GCM: self.key_template = aead.aead_key_templates.AES256_GCM_RAW elif dek_format is DekAlgorithm.AES256_SIV: # Construct AES256_SIV_RAW since it doesn't exist in Tink key_format = aes_siv_pb2.AesSivKeyFormat( # Generate 2 256-bit keys key_size=64, ) self.key_template = tink_pb2.KeyTemplate( type_url=daead.deterministic_aead_key_templates.AES256_SIV.type_url, output_prefix_type=tink_pb2.RAW, value=key_format.SerializeToString(), ) else: raise RuleError("invalid dek algorithm") def generate_key(self) -> bytes: key_data = self.registry.new_key_data(self.key_template) return key_data.value def encrypt(self, dek: bytes, plaintext: bytes, associated_data: bytes) -> bytes: key_data = tink_pb2.KeyData( type_url=self.key_template.type_url, value=dek, key_material_type=tink_pb2.KeyData.SYMMETRIC ) if self.is_deterministic: primitive = self.registry.primitive(key_data, daead.DeterministicAead) return primitive.encrypt_deterministically(plaintext, associated_data) else: primitive = self.registry.primitive(key_data, aead.Aead) return primitive.encrypt(plaintext, associated_data) def decrypt(self, dek: bytes, ciphertext: bytes, associated_data: bytes) -> bytes: key_data = tink_pb2.KeyData( type_url=self.key_template.type_url, value=dek, key_material_type=tink_pb2.KeyData.SYMMETRIC ) if self.is_deterministic: primitive = self.registry.primitive(key_data, daead.DeterministicAead) return primitive.decrypt_deterministically(ciphertext, associated_data) else: primitive = self.registry.primitive(key_data, aead.Aead) return primitive.decrypt(ciphertext, associated_data) class EncryptionExecutorTransform(object): def __init__(self, executor: EncryptionExecutor, cryptor: Cryptor, kek_name: str, dek_expiry_days: int): self._executor = executor self._cryptor = cryptor self._kek_name = kek_name self._kek = None self._dek_expiry_days = dek_expiry_days def _is_dek_rotated(self): return self._dek_expiry_days > 0 def _get_kek(self, ctx: RuleContext) -> Kek: if self._kek is None: self._kek = self._get_or_create_kek(ctx) return self._kek def _get_or_create_kek(self, ctx: RuleContext) -> Kek: is_read = ctx.rule_mode == RuleMode.READ kms_type = ctx.get_parameter(ENCRYPT_KMS_TYPE) kms_key_id = ctx.get_parameter(ENCRYPT_KMS_KEY_ID) kek_id = KekId(self._kek_name, False) kek = self._retrieve_kek_from_registry(kek_id) if kek is None: if is_read: raise RuleError(f"no kek found for {self._kek_name} during consume") if not kms_type: raise RuleError(f"no kms type found for {self._kek_name} during produce") if not kms_key_id: raise RuleError(f"no kms key id found for {self._kek_name} during produce") kek = self._store_kek_to_registry(kek_id, kms_type, kms_key_id, False) if kek is None: # handle conflicts (409) kek = self._retrieve_kek_from_registry(kek_id) if kek is None: raise RuleError(f"no kek found for {self._kek_name} during produce") if kms_type and kek.kms_type != kms_type: raise RuleError(f"found {self._kek_name} with kms type {kek.kms_type} " f"which differs from rule kms type {kms_type}") if kms_key_id and kek.kms_key_id != kms_key_id: raise RuleError(f"found {self._kek_name} with kms key id {kek.kms_key_id} " f"which differs from rule kms key id {kms_key_id}") return kek def _retrieve_kek_from_registry(self, kek_id: KekId) -> Optional[Kek]: try: return self._executor.client.get_kek(kek_id.name, kek_id.deleted) except Exception as e: if isinstance(e, SchemaRegistryError) and e.http_status_code == 404: return None raise RuleError(f"could not get kek {kek_id.name}") from e def _store_kek_to_registry( self, kek_id: KekId, kms_type: str, kms_key_id: str, shared: bool ) -> Optional[Kek]: try: return self._executor.client.register_kek(kek_id.name, kms_type, kms_key_id, shared) except Exception as e: if isinstance(e, SchemaRegistryError) and e.http_status_code == 409: return None raise RuleError(f"could not register kek {kek_id.name}") from e def _get_or_create_dek(self, ctx: RuleContext, version: Optional[int]) -> Dek: kek = self._get_kek(ctx) is_read = ctx.rule_mode == RuleMode.READ if version is None or version == 0: version = 1 dek_id = DekId( kek.name, ctx.subject, version, self._cryptor.dek_format, is_read ) dek = self._retrieve_dek_from_registry(dek_id) is_expired = self._is_expired(ctx, dek) primitive = None if dek is None or is_expired: if is_read: raise RuleError(f"no dek found for {dek_id.kek_name} during consume") encrypted_dek = None if not kek.shared: primitive = AeadWrapper(self._executor.config, self._kek) raw_dek = self._cryptor.generate_key() encrypted_dek = primitive.encrypt(raw_dek, self._cryptor.EMPTY_AAD) new_version = dek.version + 1 if is_expired else 1 try: dek = self._create_dek(dek_id, new_version, encrypted_dek) except RuleError as e: if dek is None: raise e log.warning("failed to create dek for %s, subject %s, version %d, using existing dek", kek.name, ctx.subject, new_version) key_bytes = dek.get_key_material_bytes() if key_bytes is None: if primitive is None: primitive = AeadWrapper(self._executor.config, self._kek) encrypted_dek = dek.get_encrypted_key_material_bytes() raw_dek = primitive.decrypt(encrypted_dek, self._cryptor.EMPTY_AAD) dek.set_key_material(raw_dek) return dek def _create_dek(self, dek_id: DekId, new_version: Optional[int], encrypted_dek: Optional[bytes]) -> Dek: new_dek_id = DekId( dek_id.kek_name, dek_id.subject, new_version, dek_id.algorithm, dek_id.deleted, ) dek = self._store_dek_to_registry(new_dek_id, encrypted_dek) if dek is None: # handle conflicts (409) dek = self._retrieve_dek_from_registry(dek_id) if dek is None: raise RuleError(f"no dek found for {dek_id.kek_name} during produce") return dek def _retrieve_dek_from_registry(self, key: DekId) -> Optional[Dek]: try: version = key.version if not version: version = 1 dek = self._executor.client.get_dek( key.kek_name, key.subject, key.algorithm, version, key.deleted) return dek if dek and dek.encrypted_key_material else None except Exception as e: if isinstance(e, SchemaRegistryError) and e.http_status_code == 404: return None raise RuleError(f"could not get dek for kek {key.kek_name}, subject {key.subject}") from e def _store_dek_to_registry(self, key: DekId, encrypted_dek: Optional[bytes]) -> Optional[Dek]: try: encrypted_dek_str = base64.b64encode(encrypted_dek).decode("utf-8") if encrypted_dek else None dek = self._executor.client.register_dek( key.kek_name, key.subject, encrypted_dek_str, key.algorithm, key.version) return dek except Exception as e: if isinstance(e, SchemaRegistryError) and e.http_status_code == 409: return None raise RuleError(f"could not register dek for kek {key.kek_name}, subject {key.subject}") from e def _is_expired(self, ctx: RuleContext, dek: Optional[Dek]) -> bool: now = self._executor.clock.now() return (ctx.rule_mode != RuleMode.READ and self._dek_expiry_days > 0 and dek is not None and (now - dek.ts) / MILLIS_IN_DAY > self._dek_expiry_days) def transform(self, ctx: RuleContext, field_type: FieldType, field_value: Any) -> Any: if field_value is None: return None if ctx.rule_mode == RuleMode.WRITE: plaintext = self._to_bytes(field_type, field_value) if plaintext is None: raise RuleError(f"type {field_type} not supported for encryption") version = None if self._is_dek_rotated(): version = -1 dek = self._get_or_create_dek(ctx, version) key_material_bytes = dek.get_key_material_bytes() ciphertext = self._cryptor.encrypt(key_material_bytes, plaintext, Cryptor.EMPTY_AAD) if self._is_dek_rotated(): ciphertext = self._prefix_version(dek.version, ciphertext) if field_type == FieldType.STRING: return base64.b64encode(ciphertext).decode("utf-8") else: return self._to_object(field_type, ciphertext) elif ctx.rule_mode == RuleMode.READ: if field_type == FieldType.STRING: ciphertext = base64.b64decode(field_value) else: ciphertext = self._to_bytes(field_type, field_value) if ciphertext is None: return field_value version = None if self._is_dek_rotated(): version, ciphertext = self._extract_version(ciphertext) if version is None: raise RuleError("no version found in ciphertext") dek = self._get_or_create_dek(ctx, version) key_material_bytes = dek.get_key_material_bytes() plaintext = self._cryptor.decrypt(key_material_bytes, ciphertext, Cryptor.EMPTY_AAD) return self._to_object(field_type, plaintext) else: raise RuleError(f"unsupported rule mode {ctx.rule_mode}") def _prefix_version(self, version: int, ciphertext: bytes) -> bytes: return bytes([_MAGIC_BYTE_V0]) + version.to_bytes(4, byteorder="big") + ciphertext def _extract_version(self, ciphertext: bytes) -> Tuple[Optional[int], bytes]: if len(ciphertext) < 5: return None, ciphertext version = int.from_bytes(ciphertext[1:5], byteorder="big") return version, ciphertext[5:] def _to_bytes(self, field_type: FieldType, value: Any) -> Optional[bytes]: if field_type == FieldType.STRING: return value.encode("utf-8") elif field_type == FieldType.BYTES: if isinstance(value, io.BytesIO): return value.read() return value return None def _to_object(self, field_type: FieldType, value: bytes) -> Any: if field_type == FieldType.STRING: return value.decode("utf-8") elif field_type == FieldType.BYTES: return value return None class AeadWrapper(aead.Aead): def __init__(self, config: dict, kek: Kek): self._config = config self._kek = kek self._kms_key_ids = self._get_kms_key_ids() def encrypt(self, plaintext: bytes, associated_data: bytes) -> bytes: for index, kms_key_id in enumerate(self._kms_key_ids): try: aead = self._get_aead(self._config, self._kek.kms_type, kms_key_id) return aead.encrypt(plaintext, associated_data) except Exception as e: log.warning("failed to encrypt with kek %s and kms key id %s", self._kek.name, kms_key_id) if index == len(self._kms_key_ids) - 1: raise RuleError(f"failed to encrypt with all KEKs for {self._kek.name}") from e raise RuleError("No KEK found for encryption") def decrypt(self, ciphertext: bytes, associated_data: bytes) -> bytes: for index, kms_key_id in enumerate(self._kms_key_ids): try: aead = self._get_aead(self._config, self._kek.kms_type, kms_key_id) return aead.decrypt(ciphertext, associated_data) except Exception as e: log.warning("failed to decrypt with kek %s and kms key id %s", self._kek.name, kms_key_id) if index == len(self._kms_key_ids) - 1: raise RuleError(f"failed to decrypt with all KEKs for {self._kek.name}") from e raise RuleError("No KEK found for decryption") def _get_kms_key_ids(self) -> List[str]: kms_key_ids = [self._kek.kms_key_id] alternate_kms_key_ids = None if self._kek.kms_props is not None: alternate_kms_key_ids = self._kek.kms_props.properties.get(ENCRYPT_ALTERNATE_KMS_KEY_IDS) if alternate_kms_key_ids is None: alternate_kms_key_ids = self._config.get(ENCRYPT_ALTERNATE_KMS_KEY_IDS) if alternate_kms_key_ids is not None: # Split the comma-separated list of alternate KMS key IDs and append to kms_key_ids kms_key_ids.extend([id.strip() for id in alternate_kms_key_ids.split(',') if id.strip()]) return kms_key_ids def _get_aead(self, config: dict, kms_type: str, kms_key_id: str) -> aead.Aead: kek_url = kms_type + "://" + kms_key_id kms_client = self._get_kms_client(config, kek_url) return kms_client.get_aead(kek_url) def _get_kms_client(self, config: dict, kek_url: str) -> KmsClient: driver = get_kms_driver(kek_url) try: client = kms_client_from_uri(kek_url) except TinkError: client = self._register_kms_client(driver, config, kek_url) return client def _register_kms_client(self, kms_driver: KmsDriver, config: dict, kek_url: str) -> KmsClient: kms_client = kms_driver.new_kms_client(config, kek_url) register_kms_client(kms_client) return kms_client class FieldEncryptionExecutor(FieldRuleExecutor): def __init__(self, clock: Clock = Clock()): self.executor = EncryptionExecutor(clock) def configure(self, client_conf: dict, rule_conf: dict): self.executor.configure(client_conf, rule_conf) def type(self) -> str: return "ENCRYPT" def new_transform(self, ctx: RuleContext) -> FieldTransform: executor_transform = self.executor.new_transform(ctx) transform = FieldEncryptionExecutorTransform(executor_transform) return transform.transform def close(self): if self.client is not None: self.client.__exit__() @classmethod def register(cls): RuleRegistry.register_rule_executor(FieldEncryptionExecutor()) @classmethod def register_with_clock(cls, clock: Clock) -> 'FieldEncryptionExecutor': executor = FieldEncryptionExecutor(clock) RuleRegistry.register_rule_executor(executor) return executor class FieldEncryptionExecutorTransform(object): def __init__(self, executor_transform: 'EncryptionExecutorTransform'): self.executor_transform = executor_transform def transform(self, ctx: RuleContext, field_ctx: FieldContext, field_value: Any) -> Any: return self.executor_transform.transform(ctx, field_ctx.field_type, field_value)