-
Notifications
You must be signed in to change notification settings - Fork 345
feat: Add support for Agent Identity bound tokens #1821
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
6391ca5
4b628b9
b8c8990
58dea20
a05f382
3b0ef85
f167834
a636f04
2be3904
6cc7fa8
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,243 @@ | ||
| # Copyright 2025 Google LLC | ||
| # | ||
| # 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. | ||
|
|
||
| """Helpers for Agent Identity credentials.""" | ||
|
|
||
| import base64 | ||
| import hashlib | ||
| import logging | ||
| import os | ||
| import re | ||
| import time | ||
| from urllib.parse import urlparse | ||
|
|
||
| from google.auth import environment_vars | ||
| from google.auth import exceptions | ||
|
|
||
|
|
||
| _LOGGER = logging.getLogger(__name__) | ||
|
|
||
| CRYPTOGRAPHY_NOT_FOUND_ERROR = ( | ||
| "The cryptography library is required for certificate-based authentication." | ||
| "Please install it with `pip install google-auth[cryptography]`." | ||
| ) | ||
|
|
||
| # SPIFFE trust domain patterns for Agent Identities. | ||
| _AGENT_IDENTITY_SPIFFE_TRUST_DOMAIN_PATTERNS = [ | ||
| r"^agents\.global\.org-\d+\.system\.id\.goog$", | ||
| r"^agents\.global\.proj-\d+\.system\.id\.goog$", | ||
| ] | ||
|
|
||
| # Constants for polling the certificate file. | ||
| _FAST_POLL_CYCLES = 50 | ||
| _FAST_POLL_INTERVAL = 0.1 # 100ms | ||
| _SLOW_POLL_INTERVAL = 0.5 # 500ms | ||
| _TOTAL_TIMEOUT = 30 # seconds | ||
|
|
||
| # Calculate the number of slow poll cycles based on the total timeout. | ||
| _SLOW_POLL_CYCLES = int( | ||
| (_TOTAL_TIMEOUT - (_FAST_POLL_CYCLES * _FAST_POLL_INTERVAL)) / _SLOW_POLL_INTERVAL | ||
| ) | ||
|
|
||
| _POLLING_INTERVALS = ([_FAST_POLL_INTERVAL] * _FAST_POLL_CYCLES) + ( | ||
| [_SLOW_POLL_INTERVAL] * _SLOW_POLL_CYCLES | ||
| ) | ||
|
|
||
|
|
||
| def get_agent_identity_certificate_path(): | ||
| """Gets the certificate path from the certificate config file. | ||
|
|
||
| The path to the certificate config file is read from the | ||
| GOOGLE_API_CERTIFICATE_CONFIG environment variable. This function | ||
| implements a retry mechanism to handle cases where the environment | ||
| variable is set before the files are available on the filesystem. | ||
|
|
||
| Returns: | ||
| str: The path to the leaf certificate file. | ||
|
|
||
| Raises: | ||
| google.auth.exceptions.RefreshError: If the certificate config file | ||
| or the certificate file cannot be found after retries. | ||
| """ | ||
| import json | ||
|
|
||
| cert_config_path = os.environ.get(environment_vars.GOOGLE_API_CERTIFICATE_CONFIG) | ||
| if not cert_config_path: | ||
| return None | ||
|
|
||
| has_logged_warning = False | ||
|
|
||
| for interval in _POLLING_INTERVALS: | ||
| try: | ||
| with open(cert_config_path, "r") as f: | ||
| cert_config = json.load(f) | ||
| cert_path = ( | ||
| cert_config.get("cert_configs", {}) | ||
| .get("workload", {}) | ||
| .get("cert_path") | ||
| ) | ||
| if cert_path and os.path.exists(cert_path): | ||
| return cert_path | ||
| except (IOError, ValueError, KeyError): | ||
| if not has_logged_warning: | ||
| _LOGGER.warning( | ||
| "Certificate config file not found at %s (from %s environment " | ||
| "variable). Retrying for up to %s seconds.", | ||
| cert_config_path, | ||
| environment_vars.GOOGLE_API_CERTIFICATE_CONFIG, | ||
| _TOTAL_TIMEOUT, | ||
| ) | ||
| has_logged_warning = True | ||
| pass | ||
|
|
||
| # A sleep is required in two cases: | ||
| # 1. The config file is not found (the except block). | ||
| # 2. The config file is found, but the certificate is not yet available. | ||
| # In both cases, we need to poll, so we sleep on every iteration | ||
| # that doesn't return a certificate. | ||
| time.sleep(interval) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This doesn't matter at all, but I think the sleep could be inside the except, to remove that pass
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks for the suggestion! I've kept the time.sleep() outside the except block to handle two
In both scenarios, we need to pause and retry. If we move the sleep into the except block, there won't be any pause and sleep for the second scenario would . I've added a comment to the code to make this polling logic clearer for future developers. |
||
|
|
||
| raise exceptions.RefreshError( | ||
| "Certificate config or certificate file not found after multiple retries. " | ||
lsirac marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| f"Token binding protection is failing. You can turn off this protection by setting " | ||
| f"{environment_vars.GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES} to false " | ||
| "to fall back to unbound tokens." | ||
| ) | ||
|
|
||
|
|
||
| def get_and_parse_agent_identity_certificate(): | ||
| """Gets and parses the agent identity certificate if not opted out. | ||
|
|
||
| Checks if the user has opted out of certificate-bound tokens. If not, | ||
| it gets the certificate path, reads the file, and parses it. | ||
|
|
||
| Returns: | ||
| The parsed certificate object if found and not opted out, otherwise None. | ||
| """ | ||
| # If the user has opted out of cert bound tokens, there is no need to | ||
| # look up the certificate. | ||
| is_opted_out = ( | ||
| os.environ.get( | ||
| environment_vars.GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES, | ||
| "true", | ||
| ).lower() | ||
| == "false" | ||
| ) | ||
| if is_opted_out: | ||
| return None | ||
|
|
||
| cert_path = get_agent_identity_certificate_path() | ||
| if not cert_path: | ||
| return None | ||
|
|
||
| with open(cert_path, "rb") as cert_file: | ||
| cert_bytes = cert_file.read() | ||
|
|
||
| return parse_certificate(cert_bytes) | ||
|
|
||
|
|
||
| def parse_certificate(cert_bytes): | ||
| """Parses a PEM-encoded certificate. | ||
|
|
||
| Args: | ||
| cert_bytes (bytes): The PEM-encoded certificate bytes. | ||
|
|
||
| Returns: | ||
| cryptography.x509.Certificate: The parsed certificate object. | ||
| """ | ||
| try: | ||
| from cryptography import x509 | ||
|
|
||
| return x509.load_pem_x509_certificate(cert_bytes) | ||
| except ImportError as e: | ||
| raise ImportError(CRYPTOGRAPHY_NOT_FOUND_ERROR) from e | ||
|
|
||
|
|
||
| def _is_agent_identity_certificate(cert): | ||
| """Checks if a certificate is an Agent Identity certificate. | ||
|
|
||
| This is determined by checking the Subject Alternative Name (SAN) for a | ||
| SPIFFE ID with a trust domain matching Agent Identity patterns. | ||
|
|
||
| Args: | ||
| cert (cryptography.x509.Certificate): The parsed certificate object. | ||
|
|
||
| Returns: | ||
| bool: True if the certificate is an Agent Identity certificate, | ||
| False otherwise. | ||
| """ | ||
| try: | ||
| from cryptography import x509 | ||
| from cryptography.x509.oid import ExtensionOID | ||
|
|
||
| try: | ||
| ext = cert.extensions.get_extension_for_oid( | ||
| ExtensionOID.SUBJECT_ALTERNATIVE_NAME | ||
| ) | ||
| except x509.ExtensionNotFound: | ||
| return False | ||
| uris = ext.value.get_values_for_type(x509.UniformResourceIdentifier) | ||
|
|
||
| for uri in uris: | ||
| parsed_uri = urlparse(uri) | ||
| if parsed_uri.scheme == "spiffe": | ||
| trust_domain = parsed_uri.netloc | ||
| for pattern in _AGENT_IDENTITY_SPIFFE_TRUST_DOMAIN_PATTERNS: | ||
| if re.match(pattern, trust_domain): | ||
| return True | ||
| return False | ||
| except ImportError as e: | ||
| raise ImportError(CRYPTOGRAPHY_NOT_FOUND_ERROR) from e | ||
|
|
||
|
|
||
| def calculate_certificate_fingerprint(cert): | ||
| """Calculates the base64-encoded SHA256 hash of a DER-encoded certificate. | ||
|
|
||
| Args: | ||
| cert (cryptography.x509.Certificate): The parsed certificate object. | ||
|
|
||
| Returns: | ||
| str: The base64-encoded SHA256 fingerprint. | ||
| """ | ||
| try: | ||
| from cryptography.hazmat.primitives import serialization | ||
|
|
||
| der_cert = cert.public_bytes(serialization.Encoding.DER) | ||
| fingerprint = hashlib.sha256(der_cert).digest() | ||
| return base64.urlsafe_b64encode(fingerprint).rstrip(b"=").decode("utf-8") | ||
| except ImportError as e: | ||
| raise ImportError(CRYPTOGRAPHY_NOT_FOUND_ERROR) from e | ||
|
|
||
|
|
||
| def should_request_bound_token(cert): | ||
| """Determines if a bound token should be requested. | ||
|
|
||
| This is based on the GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES | ||
| environment variable and whether the certificate is an agent identity cert. | ||
|
|
||
| Args: | ||
| cert (cryptography.x509.Certificate): The parsed certificate object. | ||
|
|
||
| Returns: | ||
| bool: True if a bound token should be requested, False otherwise. | ||
| """ | ||
| is_agent_cert = _is_agent_identity_certificate(cert) | ||
| is_opted_in = ( | ||
| os.environ.get( | ||
| environment_vars.GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES, | ||
| "true", | ||
| ).lower() | ||
| == "true" | ||
| ) | ||
| return is_agent_cert and is_opted_in | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,3 +1,3 @@ | ||
| [mypy] | ||
| python_version = 3.7 | ||
| python_version = 3.9 | ||
| namespace_packages = True |
Uh oh!
There was an error while loading. Please reload this page.