ksobrenat32 commented on code in PR #35524: URL: https://github.com/apache/beam/pull/35524#discussion_r2260876685
########## infra/keys/service_account.py: ########## @@ -0,0 +1,485 @@ +# 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 logging +import json +import time +from typing import List,Optional +from google.cloud import iam_admin_v1 +from google.cloud.iam_admin_v1 import types +from google.oauth2 import service_account +from google.auth.transport.requests import Request +from google.api_core import exceptions + +class ServiceAccountManagerLoggerAdapter(logging.LoggerAdapter): + """Logger adapter that adds a prefix to all log messages.""" + + def process(self, msg, kwargs): + return f"[ServiceAccountManager] {msg}", kwargs + +class ServiceAccountManager: + def __init__(self, project_id: str, logger: logging.Logger, max_retries: int = 3) -> None: + self.project_id = project_id + self.client = iam_admin_v1.IAMClient() + self.logger = ServiceAccountManagerLoggerAdapter(logger, {}) + self.max_retries = max_retries + self.logger.info(f"Initialized ServiceAccountManager for project: {self.project_id}") + + def _normalize_account_email(self, account_id: str) -> str: + """ + Normalizes the account identifier to a full email format. + + Args: + account_id (str): The unique identifier or email of the service account. + + Returns: + str: The full service account email address. + """ + # Handle both account ID and full email formats + if "@" in account_id and account_id.endswith(".iam.gserviceaccount.com"): + # account_id is already a full email + return account_id + else: + # account_id is just the account name + return f"{account_id}@{self.project_id}.iam.gserviceaccount.com" + + def _get_service_accounts(self) -> List[iam_admin_v1.ServiceAccount]: + """ + Retrieves all service accounts in the specified project. + + Returns: + List[iam_admin_v1.ServiceAccount]: A list of service account objects. + """ + request = types.ListServiceAccountsRequest() + request.name = f"projects/{self.project_id}" + + accounts = self.client.list_service_accounts(request=request) + self.logger.debug(f"Listed service accounts: {[account.email for account in accounts.accounts]}") + return list(accounts.accounts) + + def _service_account_exists(self, account_id: str) -> bool: + """ + Checks if a service account with the given account_id exists in the project. + + Args: + account_id (str): The unique identifier or email of the service account. + + Returns: + bool: True if the service account exists, False otherwise. + """ + try: + self.get_service_account(account_id) + return True + except exceptions.NotFound: + return False + + def _service_account_is_enabled(self, account_id: str) -> bool: + """ + Checks if a service account is enabled. + + Args: + account_id (str): The unique identifier or email of the service account. + + Returns: + bool: True if the service account is enabled, False otherwise. + """ + try: + service_account = self.get_service_account(account_id) + return not service_account.disabled + except exceptions.NotFound: + self.logger.error(f"Service account {account_id} not found") + return False + + def create_service_account(self, account_id: str, display_name: Optional[str] = None) -> types.ServiceAccount: + """ + Creates a service account in the specified project. + If the service account already exists, returns the existing account (idempotent operation). + + Args: + account_id (str): The unique identifier for the service account. + display_name (Optional[str]): A human-readable name for the service account. + Returns: + types.ServiceAccount: The created or existing service account object. + """ + request = types.CreateServiceAccountRequest() + request.account_id = account_id + request.name = f"projects/{self.project_id}" + + service_account = types.ServiceAccount() + service_account.display_name = display_name or account_id + request.service_account = service_account + + try: + account = self.client.create_service_account(request=request) + + # Wait for the service account to be created + delay = 1 + for _ in range(self.max_retries): + if self._service_account_exists(account_id): + break + time.sleep(delay) + delay *= 2 + else: + self.logger.error(f"Service account {account_id} creation timed out after {self.max_retries} retries.") + raise exceptions.DeadlineExceeded(f"Service account {account_id} creation timed out.") + + self.logger.info(f"Created service account: {account.email}") + return account + except exceptions.Conflict: + existing_account = self.get_service_account(account_id) + self.logger.info(f"Service account already exists: {existing_account.email}") + return existing_account + + def get_service_account(self, account_id: str) -> types.ServiceAccount: + """ + Retrieves a service account by its unique identifier or email. + + Args: + account_id (str): The unique identifier or email of the service account. + + Returns: + types.ServiceAccount: The service account object. + """ + service_account_email = self._normalize_account_email(account_id) + + request = types.GetServiceAccountRequest() + request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}" + + try: + service_account = self.client.get_service_account(request=request) + self.logger.info(f"Retrieved service account: {service_account.email}") + return service_account + except exceptions.NotFound: + self.logger.error(f"Service account {account_id} not found") + raise + + def enable_service_account(self, account_id: str) -> None: + """ + Enables a service account in the specified project. + + Args: + account_id (str): The unique identifier or email of the service account to enable. + """ + service_account_email = self._normalize_account_email(account_id) + request = types.EnableServiceAccountRequest() + request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}" + + self.client.enable_service_account(request=request) + + # Wait for the service account to be enabled + delay = 1 + for _ in range(self.max_retries): + if self._service_account_is_enabled(account_id): + break + time.sleep(delay) + delay *= 2 + else: + self.logger.error(f"Service account {account_id} enabling timed out after {self.max_retries} retries.") + raise exceptions.DeadlineExceeded(f"Service account {account_id} enabling timed out.") + + self.logger.info(f"Enabled service account: {account_id}") + + def disable_service_account(self, account_id: str) -> None: + """ + Disables a service account in the specified project. + + Args: + account_id (str): The unique identifier or email of the service account to disable. + """ + service_account_email = self._normalize_account_email(account_id) + request = types.DisableServiceAccountRequest() + request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}" + + self.client.disable_service_account(request=request) + + # Wait for the service account to be disabled + delay = 1 + for _ in range(self.max_retries): + if not self._service_account_is_enabled(account_id): + break + time.sleep(delay) + delay *= 2 + else: + self.logger.error(f"Service account {account_id} disabling timed out after {self.max_retries} retries.") + raise exceptions.DeadlineExceeded(f"Service account {account_id} disabling timed out.") + + self.logger.info(f"Disabled service account: {account_id}") + + def delete_service_account(self, account_id: str) -> None: + """ + Deletes a service account in the specified project. + + Args: + account_id (str): The unique identifier or email of the service account to delete. + """ + service_account_email = self._normalize_account_email(account_id) + request = types.DeleteServiceAccountRequest() + request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}" + + self.client.delete_service_account(request=request) + + # Wait for the service account to be deleted + delay = 1 + for _ in range(self.max_retries): + if not self._service_account_exists(account_id): + break + time.sleep(delay) + delay *= 2 + else: + self.logger.error(f"Service account {account_id} deletion timed out after {self.max_retries} retries.") + raise exceptions.DeadlineExceeded(f"Service account {account_id} deletion timed out.") + + self.logger.info(f"Deleted service account: {account_id}") + + def _get_service_account_keys(self, account_id: str) -> List[iam_admin_v1.ServiceAccountKey]: + """ + Retrieves all keys for the specified service account. + + Args: + account_id (str): The unique identifier or email of the service account. + + Returns: + List[iam_admin_v1.ServiceAccountKey]: A list of service account key objects. + """ + service_account_email = self._normalize_account_email(account_id) + request = types.ListServiceAccountKeysRequest() + request.name = f"projects/{self.project_id}/serviceAccounts/{service_account_email}" + + response = self.client.list_service_account_keys(request=request) + self.logger.debug(f"Listed keys for service account: {account_id}") + return list(response.keys) + + def _service_account_key_exists(self, account_id: str, key_id: str) -> bool: + """ + Checks if a service account key exists for the specified service account. + + Args: + account_id (str): The unique identifier or email of the service account. + key_id (str): The ID of the service account key to check. + + Returns: + bool: True if the key exists, False otherwise. + """ + keys = self._get_service_account_keys(account_id) + return any(key.name.endswith(key_id) for key in keys) Review Comment: One thing is the key and other the key_id, the key_id is the last part of the key name -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
