diff --git a/sdks/python/apache_beam/transforms/util.py b/sdks/python/apache_beam/transforms/util.py index 2ea9df9399cb..30f8dc01c458 100644 --- a/sdks/python/apache_beam/transforms/util.py +++ b/sdks/python/apache_beam/transforms/util.py @@ -82,6 +82,9 @@ from apache_beam.utils import shared from apache_beam.utils import windowed_value from apache_beam.utils.annotations import deprecated +from apache_beam.utils.secret import GcpHsmGeneratedSecret +from apache_beam.utils.secret import GcpSecret +from apache_beam.utils.secret import Secret from apache_beam.utils.sharded_key import ShardedKey from apache_beam.utils.timestamp import Timestamp @@ -94,6 +97,7 @@ 'BatchElements', 'CoGroupByKey', 'Distinct', + 'GcpHsmGeneratedSecret', 'GcpSecret', 'GroupByEncryptedKey', 'Keys', @@ -327,247 +331,6 @@ def RemoveDuplicates(pcoll): return pcoll | 'RemoveDuplicates' >> Distinct() -class Secret(): - """A secret management class used for handling sensitive data. - - This class provides a generic interface for secret management. Implementations - of this class should handle fetching secrets from a secret management system. - """ - def get_secret_bytes(self) -> bytes: - """Returns the secret as a byte string.""" - raise NotImplementedError() - - @staticmethod - def generate_secret_bytes() -> bytes: - """Generates a new secret key.""" - return Fernet.generate_key() - - @staticmethod - def parse_secret_option(secret) -> 'Secret': - """Parses a secret string and returns the appropriate secret type. - - The secret string should be formatted like: - 'type:;:' - - For example, 'type:GcpSecret;version_name:my_secret/versions/latest' - would return a GcpSecret initialized with 'my_secret/versions/latest'. - """ - param_map = {} - for param in secret.split(';'): - parts = param.split(':') - param_map[parts[0]] = parts[1] - - if 'type' not in param_map: - raise ValueError('Secret string must contain a valid type parameter') - - secret_type = param_map['type'].lower() - del param_map['type'] - secret_class = Secret - secret_params = None - if secret_type == 'gcpsecret': - secret_class = GcpSecret # type: ignore[assignment] - secret_params = ['version_name'] - elif secret_type == 'gcphsmgeneratedsecret': - secret_class = GcpHsmGeneratedSecret # type: ignore[assignment] - secret_params = [ - 'project_id', 'location_id', 'key_ring_id', 'key_id', 'job_name' - ] - else: - raise ValueError( - f'Invalid secret type {secret_type}, currently only ' - 'GcpSecret and GcpHsmGeneratedSecret are supported') - - for param_name in param_map.keys(): - if param_name not in secret_params: - raise ValueError( - f'Invalid secret parameter {param_name}, ' - f'{secret_type} only supports the following ' - f'parameters: {secret_params}') - return secret_class(**param_map) - - -class GcpSecret(Secret): - """A secret manager implementation that retrieves secrets from Google Cloud - Secret Manager. - """ - def __init__(self, version_name: str): - """Initializes a GcpSecret object. - - Args: - version_name: The full version name of the secret in Google Cloud Secret - Manager. For example: - projects//secrets//versions/1. - For more info, see - https://cloud.google.com/python/docs/reference/secretmanager/latest/google.cloud.secretmanager_v1beta1.services.secret_manager_service.SecretManagerServiceClient#google_cloud_secretmanager_v1beta1_services_secret_manager_service_SecretManagerServiceClient_access_secret_version - """ - self._version_name = version_name - - def get_secret_bytes(self) -> bytes: - try: - from google.cloud import secretmanager - client = secretmanager.SecretManagerServiceClient() - response = client.access_secret_version( - request={"name": self._version_name}) - secret = response.payload.data - return secret - except Exception as e: - raise RuntimeError( - 'Failed to retrieve secret bytes for secret ' - f'{self._version_name} with exception {e}') - - def __eq__(self, secret): - return self._version_name == getattr(secret, '_version_name', None) - - -class GcpHsmGeneratedSecret(Secret): - """A secret manager implementation that generates a secret using a GCP HSM key - and stores it in Google Cloud Secret Manager. If the secret already exists, - it will be retrieved. - """ - def __init__( - self, - project_id: str, - location_id: str, - key_ring_id: str, - key_id: str, - job_name: str): - """Initializes a GcpHsmGeneratedSecret object. - - Args: - project_id: The GCP project ID. - location_id: The GCP location ID for the HSM key. - key_ring_id: The ID of the KMS key ring. - key_id: The ID of the KMS key. - job_name: The name of the job, used to generate a unique secret name. - """ - self._project_id = project_id - self._location_id = location_id - self._key_ring_id = key_ring_id - self._key_id = key_id - self._secret_version_name = f'HsmGeneratedSecret_{job_name}' - - def get_secret_bytes(self) -> bytes: - """Retrieves the secret bytes. - - If the secret version already exists in Secret Manager, it is retrieved. - Otherwise, a new secret and version are created. The new secret is - generated using the HSM key. - - Returns: - The secret as a byte string. - """ - try: - from google.api_core import exceptions as api_exceptions - from google.cloud import secretmanager - client = secretmanager.SecretManagerServiceClient() - - project_path = f"projects/{self._project_id}" - secret_path = f"{project_path}/secrets/{self._secret_version_name}" - # Since we may generate multiple versions when doing this on workers, - # just always take the first version added to maintain consistency. - secret_version_path = f"{secret_path}/versions/1" - - try: - response = client.access_secret_version( - request={"name": secret_version_path}) - return response.payload.data - except api_exceptions.NotFound: - # Don't bother logging yet, we'll only log if we actually add the - # secret version below - pass - - try: - client.create_secret( - request={ - "parent": project_path, - "secret_id": self._secret_version_name, - "secret": { - "replication": { - "automatic": {} - } - }, - }) - except api_exceptions.AlreadyExists: - # Don't bother logging yet, we'll only log if we actually add the - # secret version below - pass - - new_key = self.generate_dek() - try: - # Try one more time in case it was created while we were generating the - # DEK. - response = client.access_secret_version( - request={"name": secret_version_path}) - return response.payload.data - except api_exceptions.NotFound: - _LOGGER.info( - "Secret version %s not found. " - "Creating new secret and version.", - secret_version_path) - client.add_secret_version( - request={ - "parent": secret_path, "payload": { - "data": new_key - } - }) - response = client.access_secret_version( - request={"name": secret_version_path}) - return response.payload.data - - except Exception as e: - raise RuntimeError( - f'Failed to retrieve or create secret bytes for secret ' - f'{self._secret_version_name} with exception {e}') - - def generate_dek(self, dek_size: int = 32) -> bytes: - """Generates a new Data Encryption Key (DEK) using an HSM-backed key. - - This function follows a key derivation process that incorporates entropy - from the HSM-backed key into the nonce used for key derivation. - - Args: - dek_size: The size of the DEK to generate. - - Returns: - A new DEK of the specified size, url-safe base64-encoded. - """ - try: - import base64 - import os - - from cryptography.hazmat.primitives import hashes - from cryptography.hazmat.primitives.kdf.hkdf import HKDF - from google.cloud import kms - - # 1. Generate a random nonce (nonce_one) - nonce_one = os.urandom(dek_size) - - # 2. Use the HSM-backed key to encrypt nonce_one to create nonce_two - kms_client = kms.KeyManagementServiceClient() - key_path = kms_client.crypto_key_path( - self._project_id, self._location_id, self._key_ring_id, self._key_id) - response = kms_client.encrypt( - request={ - 'name': key_path, 'plaintext': nonce_one - }) - nonce_two = response.ciphertext - - # 3. Generate a Derivation Key (DK) - dk = os.urandom(dek_size) - - # 4. Use a KDF to derive the DEK using DK and nonce_two - hkdf = HKDF( - algorithm=hashes.SHA256(), - length=dek_size, - salt=nonce_two, - info=None, - ) - dek = hkdf.derive(dk) - return base64.urlsafe_b64encode(dek) - except Exception as e: - raise RuntimeError(f'Failed to generate DEK with exception {e}') - - class _EncryptMessage(DoFn): """A DoFn that encrypts the key and value of each element.""" def __init__( diff --git a/sdks/python/apache_beam/transforms/util_test.py b/sdks/python/apache_beam/transforms/util_test.py index 63ce42726c1f..c5c4e63e4d04 100644 --- a/sdks/python/apache_beam/transforms/util_test.py +++ b/sdks/python/apache_beam/transforms/util_test.py @@ -243,7 +243,7 @@ def test_co_group_by_key_on_unpickled(self): assert_that(pcoll, equal_to(expected)) -class FakeSecret(beam.Secret): +class FakeSecret(beam.utils.secret.Secret): def __init__(self, version_name=None, should_throw=False): self._secret = b'aKwI2PmqYFt2p5tNKCyBS5qYmHhHsGZcyZrnZQiQ-uE=' self._should_throw = should_throw @@ -308,7 +308,8 @@ def test_secret_manager_parses_correctly(self, secret_string, secret): exception_str='must contain a valid type parameter'), param( secret_string='type:gcpsecreT', - exception_str='missing 1 required positional argument'), + exception_str= + r"Secret name \('name'\) must be specified in secret spec."), param( secret_string='type:gcpsecreT;version_name:foo;extra:val', exception_str='Invalid secret parameter extra'), @@ -387,7 +388,8 @@ def test_gbek_fake_secret_manager_actually_does_encryption(self): result, equal_to([('a', ([1, 2])), ('b', ([3])), ('c', ([4]))])) @mock.patch('apache_beam.transforms.util._DecryptMessage', MockNoOpDecrypt) - @mock.patch('apache_beam.transforms.util.GcpSecret', FakeSecret) + @mock.patch.dict( + 'apache_beam.utils.secret._SECRET_CLASSES', {'gcpsecret': FakeSecret}) def test_gbk_actually_does_encryption(self): options = PipelineOptions() # Version of GcpSecret doesn't matter since it is replaced by FakeSecret diff --git a/sdks/python/apache_beam/utils/secret.py b/sdks/python/apache_beam/utils/secret.py new file mode 100644 index 000000000000..608aba77c400 --- /dev/null +++ b/sdks/python/apache_beam/utils/secret.py @@ -0,0 +1,493 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You 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. +# + +"""Interface and implementations for Secret providers in Apache Beam.""" + +import abc +import json +import logging +import os +import warnings +from typing import Any, Dict, Optional, Union + +from apache_beam.utils.annotations import deprecated + +__all__ = [ + 'Secret', + 'RawSecret', + 'GcpSecret', + 'GcpHsmGeneratedSecret', + 'generate_secret_bytes', +] + +_LOGGER = logging.getLogger(__name__) + + +def generate_secret_bytes() -> bytes: + """Generates a new secret key using Fernet.""" + from cryptography.fernet import Fernet + return Fernet.generate_key() + + +class Secret(abc.ABC): + """A secret management class used for handling sensitive data. + + This class provides a generic interface for secret management. Implementations + of this class should handle fetching secrets from a secret management system. + """ + def __init__(self): + self._cached_secret_bytes: Optional[bytes] = None + + def get(self, cacheSecret: bool = False) -> str: + """Retrieve secret value as string. + + Args: + cacheSecret: If True, caches secret value in memory after first fetch. + + Returns: + The retrieved secret value as string. + """ + return self.get_bytes(cacheSecret=cacheSecret).decode("utf-8") + + def get_bytes(self, cacheSecret: bool = False) -> bytes: + """Retrieve secret value as bytes. + + Args: + cacheSecret: If True, caches secret value in memory after first fetch. + + Returns: + The retrieved secret value as bytes. + """ + if cacheSecret and getattr(self, '_cached_secret_bytes', None) is not None: + return self._cached_secret_bytes + + secret_val_bytes = self.get_secret_bytes() + + if cacheSecret: + self._cached_secret_bytes = secret_val_bytes + + return secret_val_bytes + + @abc.abstractmethod + def get_secret_bytes(self) -> bytes: + """Retrieve secret value as bytes from the underlying secret provider. + + Returns: + The retrieved secret value as bytes. + """ + raise NotImplementedError + + @staticmethod + @deprecated(since='2.77.0', current='generate_secret_bytes') + def generate_secret_bytes() -> bytes: + """Generates a new secret key. + + Deprecated: Use global :func:`generate_secret_bytes` instead. + """ + return generate_secret_bytes() + + @classmethod + @deprecated(since='2.77.0', current='from_option_string') + def parse_secret_option(cls, secret: str) -> 'Secret': + """Parses a secret string and returns the appropriate secret type. + + The secret string should be formatted like: + 'type:;:' + + For example, 'type:GcpSecret;version_name:my_secret/versions/latest' + would return a GcpSecret initialized with 'my_secret/versions/latest'. + + Deprecated: Use :meth:`from_option_string` instead. + """ + return cls.from_option_string(secret) + + def __getstate__(self): + """Strip cached secrets before pickling for pipeline submission/transmission.""" + state = self.__dict__.copy() + state['_cached_secret_bytes'] = None + return state + + @classmethod + def from_spec( + cls, + spec: Union[str, Dict[str, str]], + secret_manager: Optional[str] = None, + secret_type: Optional[str] = None) -> 'Secret': + """Return a Secret instance based on secret_manager provider and secret specification. + + Args: + spec: Secret string (raw secret or JSON specification string). + secret_manager: Provider type string (e.g. 'GoogleCloudSecretManager'). + secret_type: Provider type string (e.g. 'gcpsecret'). + + Returns: + An instance of Secret. + """ + sm_manager = secret_manager.strip( + ) if secret_manager and secret_manager.strip() else None + sm_type = secret_type.strip( + ) if secret_type and secret_type.strip() else None + + if sm_manager and sm_type: + raise ValueError( + f"Cannot specify both 'secret_manager' ('{secret_manager}') and 'secret_type' ('{secret_type}'). " + "Please specify only one.") + + if isinstance(spec, str): + spec_dict = None + try: + spec_dict = json.loads(spec) + if not isinstance(spec_dict, dict): + spec_dict = None + except Exception: + pass + elif isinstance(spec, dict): + spec_dict = spec + else: + spec_dict = None + + provider_str = sm_manager or sm_type + if provider_str: + secret_cls = _SECRET_CLASSES.get(provider_str.lower()) + if secret_cls: + if isinstance(spec_dict, dict) and hasattr(secret_cls, 'from_dict'): + return secret_cls.from_dict(spec_dict) + elif isinstance(spec_dict, dict): + return secret_cls(**spec_dict) + else: + return secret_cls(spec) + else: + raise ValueError( + f"Unsupported secret provider: '{provider_str}'. Currently supported options: 'GoogleCloudSecretManager', 'gcpsecret'." + ) + + # If secret_manager is not set or empty, check if spec is a JSON specification dict + if spec_dict is not None: + msg = ( + "The 'spec' parameter appears to be a JSON specification, but " + "'secret_manager' is not set. Defaulting to Raw.") + _LOGGER.warning(msg) + warnings.warn(msg, UserWarning) + + return RawSecret(spec) + + @classmethod + def from_option_string(cls, option: str) -> 'Secret': + param_map = {} + for param in option.split(';'): + parts = param.split(':') + if len(parts) == 2: + param_map[parts[0]] = parts[1] + + if 'type' not in param_map: + raise ValueError('Secret string must contain a valid type parameter') + + secret_manager = param_map.pop('type') + return cls.from_spec(json.dumps(param_map), secret_manager) + + +class RawSecret(Secret): + """Secret implementation wrapping a raw secret string or bytes directly.""" + def __init__(self, secret: Union[str, bytes]): + super().__init__() + if isinstance(secret, str): + self._secret = secret.encode("utf-8") + else: + self._secret = secret + + def get_secret_bytes(self) -> bytes: + return self._secret + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, RawSecret): + return False + return self._secret == other._secret + + +class GcpSecret(Secret): + """A secret manager implementation that retrieves secrets from Google Cloud + Secret Manager. + """ + def __init__(self, version_name: str): + """Initializes a GcpSecret object. + + Args: + version_name: The full version name of the secret in Google Cloud Secret + Manager. For example: + projects//secrets//versions/1. + For more info, see + https://cloud.google.com/python/docs/reference/secretmanager/latest/google.cloud.secretmanager_v1beta1.services.secret_manager_service.SecretManagerServiceClient#google_cloud_secretmanager_v1beta1_services_secret_manager_service_SecretManagerServiceClient_access_secret_version + """ + super().__init__() + self._version_name = version_name + + @classmethod + def from_dict(cls, spec_dict: Dict[str, str]) -> 'GcpSecret': + """Initialize GcpSecret from a dictionary specification.""" + allowed_keys = {'version_name', 'name', 'project', 'version'} + invalid_keys = set(spec_dict.keys()) - allowed_keys + if invalid_keys: + raise ValueError( + f"Invalid secret parameter {', '.join(sorted(invalid_keys))}") + version_name = cls._parse_version_name(spec_dict) + return cls(version_name) + + @classmethod + def _parse_version_name(cls, spec_dict: Dict[str, str]) -> str: + if "version_name" in spec_dict: + return spec_dict["version_name"] + + secret_id = spec_dict.get("name") + if not secret_id: + raise ValueError("Secret name ('name') must be specified in secret spec.") + + # Resolve project ID from spec, environment variables, or Application Default Credentials + project_id = ( + spec_dict.get("project") or os.environ.get("GOOGLE_CLOUD_PROJECT") or + os.environ.get("GCP_PROJECT")) + + if not project_id: + try: + import google.auth + _, project_id = google.auth.default() + except Exception: + pass + + version_id = spec_dict.get("version", "latest") + + if not project_id: + raise ValueError( + f"Could not resolve GCP project ID for secret '{secret_id}'. " + "Please specify 'project' in the secret spec, set GOOGLE_CLOUD_PROJECT environment variable, " + "or configure Application Default Credentials.") + + return f"projects/{project_id}/secrets/{secret_id}/versions/{version_id}" + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, GcpSecret): + return False + return self._version_name == other._version_name + + def get_secret_bytes(self) -> bytes: + """Get the secret value as bytes from GCP Secret Manager. + + Returns: + The secret bytes. + """ + try: + from google.cloud import secretmanager + client = secretmanager.SecretManagerServiceClient() + response = client.access_secret_version( + request={"name": self._version_name}) + secret_val_bytes = response.payload.data + _LOGGER.info( + "Successfully fetched secret from GCP Secret Manager (version_name '%s')", + self._version_name) + return secret_val_bytes + except Exception as e: + raise RuntimeError( + f'Failed to retrieve secret bytes for secret ' + f'{self._version_name} with exception {e}') + + +class GcpHsmGeneratedSecret(Secret): + """Secret manager implementation that generates a secret using a GCP HSM key + and stores it in Google Cloud Secret Manager. If the secret already exists, + it will be retrieved. + """ + def __init__( + self, + project_id: str, + location_id: str, + key_ring_id: str, + key_id: str, + job_name: str): + """Initializes a GcpHsmGeneratedSecret object. + + Args: + project_id: The GCP project ID. + location_id: The GCP location ID for the HSM key. + key_ring_id: The ID of the KMS key ring. + key_id: The ID of the KMS key. + job_name: The name of the job, used to generate a unique secret name. + """ + super().__init__() + self._project_id = project_id + self._location_id = location_id + self._key_ring_id = key_ring_id + self._key_id = key_id + self._job_name = job_name + self._secret_version_name = f'HsmGeneratedSecret_{job_name}' + + @classmethod + def from_dict(cls, spec_dict: Dict[str, str]) -> 'GcpHsmGeneratedSecret': + """Initialize GcpHsmGeneratedSecret from a dictionary specification.""" + allowed_keys = { + 'project_id', 'location_id', 'key_ring_id', 'key_id', 'job_name' + } + missing = allowed_keys - set(spec_dict.keys()) + if missing: + raise ValueError( + f"Missing required parameter(s) for GcpHsmGeneratedSecret: {sorted(list(missing))}" + ) + invalid_keys = set(spec_dict.keys()) - allowed_keys + if invalid_keys: + raise ValueError( + f"Invalid secret parameter {', '.join(sorted(invalid_keys))}") + return cls( + project_id=spec_dict['project_id'], + location_id=spec_dict['location_id'], + key_ring_id=spec_dict['key_ring_id'], + key_id=spec_dict['key_id'], + job_name=spec_dict['job_name'], + ) + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, GcpHsmGeneratedSecret): + return False + return ( + self._project_id == other._project_id and + self._location_id == other._location_id and + self._key_ring_id == other._key_ring_id and + self._key_id == other._key_id and self._job_name == other._job_name) + + def get_secret_bytes(self) -> bytes: + """Retrieves the secret bytes. + + If the secret version already exists in Secret Manager, it is retrieved. + Otherwise, a new secret and version are created. The new secret is + generated using the HSM key. + + Returns: + The secret bytes. + """ + try: + from google.api_core import exceptions as api_exceptions + from google.cloud import secretmanager + + client = secretmanager.SecretManagerServiceClient() + + project_path = f"projects/{self._project_id}" + secret_path = f"{project_path}/secrets/{self._secret_version_name}" + # Since we may generate multiple versions when doing this on workers, + # just always take the first version added to maintain consistency. + secret_version_path = f"{secret_path}/versions/1" + + try: + response = client.access_secret_version( + request={"name": secret_version_path}) + return response.payload.data + except api_exceptions.NotFound: + # Don't bother logging yet, we'll only log if we actually add the + # secret version below + pass + + try: + client.create_secret( + request={ + "parent": project_path, + "secret_id": self._secret_version_name, + "secret": { + "replication": { + "automatic": {} + } + }, + }) + except api_exceptions.AlreadyExists: + # Don't bother logging yet, we'll only log if we actually add the + # secret version below + pass + + new_key = self.generate_dek() + try: + # Try one more time in case it was created while we were generating the DEK. + response = client.access_secret_version( + request={"name": secret_version_path}) + return response.payload.data + except api_exceptions.NotFound: + _LOGGER.info( + "Secret version %s not found. Creating new secret and version.", + secret_version_path) + client.add_secret_version( + request={ + "parent": secret_path, "payload": { + "data": new_key + } + }) + response = client.access_secret_version( + request={"name": secret_version_path}) + return response.payload.data + + except Exception as e: + raise RuntimeError( + f'Failed to retrieve or create secret bytes for secret ' + f'{self._secret_version_name} with exception {e}') + + def generate_dek(self, dek_size: int = 32) -> bytes: + """Generates a new Data Encryption Key (DEK) using an HSM-backed key. + + This function follows a key derivation process that incorporates entropy + from the HSM-backed key into the nonce used for key derivation. + + Args: + dek_size: The size of the DEK to generate. + + Returns: + A new DEK of the specified size, url-safe base64-encoded. + """ + try: + import base64 + import os + + from cryptography.hazmat.primitives import hashes + from cryptography.hazmat.primitives.kdf.hkdf import HKDF + from google.cloud import kms + + # 1. Generate a random nonce (nonce_one) + nonce_one = os.urandom(dek_size) + + # 2. Use the HSM-backed key to encrypt nonce_one to create nonce_two + kms_client = kms.KeyManagementServiceClient() + key_path = kms_client.crypto_key_path( + self._project_id, self._location_id, self._key_ring_id, self._key_id) + response = kms_client.encrypt( + request={ + 'name': key_path, 'plaintext': nonce_one + }) + nonce_two = response.ciphertext + + # 3. Generate a Derivation Key (DK) + dk = os.urandom(dek_size) + + # 4. Use a KDF to derive the DEK using DK and nonce_two + hkdf = HKDF( + algorithm=hashes.SHA256(), + length=dek_size, + salt=nonce_two, + info=None, + ) + dek = hkdf.derive(dk) + return base64.urlsafe_b64encode(dek) + except Exception as e: + raise RuntimeError(f'Failed to generate DEK with exception {e}') + + +_SECRET_CLASSES: Dict[str, type] = { + "googlecloudsecretmanager": GcpSecret, + "gcpsecret": GcpSecret, + "gcphsmgeneratedsecret": GcpHsmGeneratedSecret, +} diff --git a/sdks/python/apache_beam/utils/secret_test.py b/sdks/python/apache_beam/utils/secret_test.py new file mode 100644 index 000000000000..e43d51f41be2 --- /dev/null +++ b/sdks/python/apache_beam/utils/secret_test.py @@ -0,0 +1,315 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You 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 json +import unittest +from unittest.mock import MagicMock, patch + +from apache_beam.utils.annotations import BeamDeprecationWarning +from apache_beam.utils.secret import GcpHsmGeneratedSecret +from apache_beam.utils.secret import GcpSecret +from apache_beam.utils.secret import RawSecret +from apache_beam.utils.secret import Secret +from apache_beam.utils.secret import generate_secret_bytes + +try: + from google.cloud import secretmanager +except ImportError: + secretmanager = None # type: ignore[assignment] + + +@unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') +class GcpSecretTest(unittest.TestCase): + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_gcp_secret_success(self, mock_client_cls): + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"secret-payload-value" + mock_client.access_secret_version.return_value = mock_response + + spec = {"name": "my-secret", "version": "1", "project": "my-project"} + secret = GcpSecret.from_dict(spec) + + secret_val = secret.get(cacheSecret=True) + self.assertEqual(secret_val, "secret-payload-value") + secret_bytes = secret.get_bytes(cacheSecret=True) + self.assertEqual(secret_bytes, b"secret-payload-value") + mock_client.access_secret_version.assert_called_once_with( + request={"name": "projects/my-project/secrets/my-secret/versions/1"}) + + # Second call with cacheSecret=True should return cached value without calling client again + mock_client.reset_mock() + secret_val_cached = secret.get(cacheSecret=True) + self.assertEqual(secret_val_cached, "secret-payload-value") + mock_client.access_secret_version.assert_not_called() + + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_gcp_secret_get_bytes_uncached(self, mock_client_cls): + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"secret-payload-value" + mock_client.access_secret_version.return_value = mock_response + + spec = {"name": "my-secret", "project": "my-project"} + secret = GcpSecret.from_dict(spec) + + secret_bytes = secret.get_bytes() + self.assertEqual(secret_bytes, b"secret-payload-value") + self.assertIsNone(secret._cached_secret_bytes) + + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_gcp_secret_getstate_clears_cached_secret(self, mock_client_cls): + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"secret-payload-value" + mock_client.access_secret_version.return_value = mock_response + + spec = {"name": "my-secret", "project": "my-project"} + secret = GcpSecret.from_dict(spec) + + # Cache the secret in memory + secret.get(cacheSecret=True) + self.assertEqual(secret._cached_secret_bytes, b"secret-payload-value") + + # When pickled / getstate is called during pipeline submission + state = secret.__getstate__() + self.assertIsNone(state["_cached_secret_bytes"]) + + @patch.dict("os.environ", {"GOOGLE_CLOUD_PROJECT": "env-project-123"}) + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_gcp_secret_env_project_fallback(self, mock_client_cls): + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"env-secret-val" + mock_client.access_secret_version.return_value = mock_response + + # Project omitted from spec + spec = {"name": "env-secret", "version": "latest"} + secret = GcpSecret.from_dict(spec) + + secret_val = secret.get(cacheSecret=False) + self.assertEqual(secret_val, "env-secret-val") + self.assertEqual(secret.get_bytes(cacheSecret=False), b"env-secret-val") + mock_client.access_secret_version.assert_called_with( + request={ + "name": "projects/env-project-123/secrets/env-secret/versions/latest" + }) + + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_gcp_secret_failure_raises_exception(self, mock_client_cls): + mock_client = MagicMock() + mock_client_cls.return_value = mock_client + mock_client.access_secret_version.side_effect = RuntimeError( + "Permission denied or secret not found") + + spec = {"name": "non-existent-secret", "project": "my-project"} + secret = GcpSecret.from_dict(spec) + + with self.assertRaises(RuntimeError) as ctx: + secret.get(cacheSecret=False) + self.assertIn("Permission denied or secret not found", str(ctx.exception)) + + @patch.dict("os.environ", {}, clear=True) + @patch("google.auth.default", side_effect=Exception("No ADC")) + def test_ill_formed_missing_project_raises_value_error( + self, mock_auth_default): + spec = {"name": "my-secret"} + with self.assertRaises(ValueError) as ctx: + GcpSecret.from_dict(spec) + self.assertIn("Could not resolve GCP project ID", str(ctx.exception)) + + def test_ill_formed_missing_secret_name_raises_value_error(self): + spec = {"project": "my-project"} + with self.assertRaises(ValueError) as ctx: + GcpSecret.from_dict(spec) + self.assertIn("Secret name ('name') must be specified", str(ctx.exception)) + + +@unittest.skipIf(secretmanager is None, 'GCP dependencies are not installed') +class GcpHsmGeneratedSecretTest(unittest.TestCase): + def test_from_dict_success(self): + spec = { + "project_id": "test-proj", + "location_id": "global", + "key_ring_id": "ring", + "key_id": "key", + "job_name": "my-job" + } + secret = GcpHsmGeneratedSecret.from_dict(spec) + self.assertEqual(secret._project_id, "test-proj") + self.assertEqual(secret._location_id, "global") + self.assertEqual(secret._key_ring_id, "ring") + self.assertEqual(secret._key_id, "key") + self.assertEqual(secret._job_name, "my-job") + self.assertEqual(secret._secret_version_name, "HsmGeneratedSecret_my-job") + + def test_from_dict_missing_params_raises_value_error(self): + spec = {"project_id": "test-proj", "location_id": "global"} + with self.assertRaises(ValueError) as ctx: + GcpHsmGeneratedSecret.from_dict(spec) + self.assertIn("Missing required parameter(s)", str(ctx.exception)) + + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_get_bytes_cached(self, mock_sm_client_cls): + mock_client = MagicMock() + mock_sm_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"hsm-derived-key" + mock_client.access_secret_version.return_value = mock_response + + secret = GcpHsmGeneratedSecret("p", "l", "r", "k", "j") + secret_bytes = secret.get_bytes(cacheSecret=True) + self.assertEqual(secret_bytes, b"hsm-derived-key") + + # Second call uses cache + mock_client.reset_mock() + self.assertEqual(secret.get_bytes(cacheSecret=True), b"hsm-derived-key") + mock_client.access_secret_version.assert_not_called() + + @patch("google.cloud.secretmanager.SecretManagerServiceClient") + def test_getstate_clears_cached_secret(self, mock_sm_client_cls): + mock_client = MagicMock() + mock_sm_client_cls.return_value = mock_client + mock_response = MagicMock() + mock_response.payload.data = b"hsm-derived-key" + mock_client.access_secret_version.return_value = mock_response + + secret = GcpHsmGeneratedSecret("p", "l", "r", "k", "j") + secret.get_bytes(cacheSecret=True) + self.assertEqual(secret._cached_secret_bytes, b"hsm-derived-key") + + state = secret.__getstate__() + self.assertIsNone(state["_cached_secret_bytes"]) + + +class RawSecretTest(unittest.TestCase): + def test_raw_secret_str(self): + secret = RawSecret("STATIC_SECRET_") + self.assertEqual(secret.get(cacheSecret=True), "STATIC_SECRET_") + self.assertEqual(secret.get_bytes(cacheSecret=True), b"STATIC_SECRET_") + + def test_raw_secret_bytes(self): + secret = RawSecret(b"STATIC_BYTES_") + self.assertEqual(secret.get(cacheSecret=True), "STATIC_BYTES_") + self.assertEqual(secret.get_bytes(cacheSecret=True), b"STATIC_BYTES_") + + +class SecretFactoryTest(unittest.TestCase): + def test_secret_factory(self): + spec = json.dumps({"name": "test-secret", "project": "proj"}) + + # When provider is set to 'GoogleCloudSecretManager' + secret_gcp = Secret.from_spec( + spec=spec, secret_manager="GoogleCloudSecretManager") + self.assertIsInstance(secret_gcp, GcpSecret) + + # When provider is set to 'gcpsecret' + secret_gcp_alias = Secret.from_spec(spec=spec, secret_manager="gcpsecret") + self.assertIsInstance(secret_gcp_alias, GcpSecret) + + # When provider is None or empty with plain string + secret_raw = Secret.from_spec(spec="STATIC_SECRET_", secret_manager=None) + self.assertIsInstance(secret_raw, RawSecret) + + # Unsupported provider raises ValueError + with self.assertRaises(ValueError): + Secret.from_spec(spec="spec", secret_manager="unsupported_provider") + + def test_secret_factory_secret_type(self): + spec = json.dumps({"name": "test-secret", "project": "proj"}) + + # Using secret_type alone + secret_type_obj = Secret.from_spec(spec=spec, secret_type="gcpsecret") + self.assertIsInstance(secret_type_obj, GcpSecret) + + def test_secret_factory_both_providers_set_raises_exception(self): + spec = json.dumps({"name": "test-secret", "project": "proj"}) + with self.assertRaises(ValueError) as ctx: + Secret.from_spec( + spec=spec, + secret_manager="GoogleCloudSecretManager", + secret_type="gcpsecret") + self.assertIn("Cannot specify both 'secret_manager'", str(ctx.exception)) + + def test_json_secret_without_secret_manager_warning(self): + json_spec = json.dumps({"name": "my-secret", "project": "my-proj"}) + with self.assertWarns(UserWarning): + secret = Secret.from_spec(spec=json_spec, secret_manager=None) + self.assertIsInstance(secret, RawSecret) + + def test_from_option_string(self): + option_str = "type:gcpsecreT;version_name:my_secret/versions/latest" + secret = Secret.from_option_string(option_str) + self.assertIsInstance(secret, GcpSecret) + self.assertEqual(secret._version_name, "my_secret/versions/latest") + + def test_from_option_string_hsm(self): + option_str = ( + "type:gcphsmgeneratedsecret;project_id:p;location_id:l;" + "key_ring_id:r;key_id:k;job_name:j") + secret = Secret.from_option_string(option_str) + self.assertIsInstance(secret, GcpHsmGeneratedSecret) + self.assertEqual(secret._project_id, "p") + self.assertEqual(secret._job_name, "j") + + def test_generate_secret_bytes(self): + key = generate_secret_bytes() + self.assertIsInstance(key, bytes) + self.assertTrue(len(key) > 0) + + def test_deprecated_parse_secret_option(self): + option_str = "type:gcpsecreT;version_name:my_secret/versions/latest" + with self.assertWarns(BeamDeprecationWarning): + secret = Secret.parse_secret_option(option_str) + self.assertIsInstance(secret, GcpSecret) + self.assertEqual(secret._version_name, "my_secret/versions/latest") + + def test_deprecated_generate_secret_bytes(self): + with self.assertWarns(BeamDeprecationWarning): + key = Secret.generate_secret_bytes() + self.assertIsInstance(key, bytes) + self.assertTrue(len(key) > 0) + + def test_equality(self): + raw1 = RawSecret("secret_value") + raw2 = RawSecret("secret_value") + raw3 = RawSecret("other_value") + self.assertEqual(raw1, raw2) + self.assertNotEqual(raw1, raw3) + self.assertNotEqual(raw1, "secret_value") + + gcp1 = GcpSecret.from_dict({"name": "sec", "project": "proj"}) + gcp2 = GcpSecret.from_dict({"name": "sec", "project": "proj"}) + gcp3 = GcpSecret.from_dict({"name": "other", "project": "proj"}) + self.assertEqual(gcp1, gcp2) + self.assertNotEqual(gcp1, gcp3) + self.assertNotEqual(gcp1, raw1) + + hsm1 = GcpHsmGeneratedSecret("p", "l", "r", "k", "j") + hsm2 = GcpHsmGeneratedSecret("p", "l", "r", "k", "j") + hsm3 = GcpHsmGeneratedSecret("p", "l", "r", "k", "other") + self.assertEqual(hsm1, hsm2) + self.assertNotEqual(hsm1, hsm3) + self.assertNotEqual(hsm1, gcp1) + + +if __name__ == "__main__": + unittest.main()