feat: Add new gcloud commands, API clients, and third-party libraries across various services.

gcloud auth application-default login/google-cloud-sdk/platform/bq/auth/gcloud_credential_loader.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python
+"""Utilities to load Google Auth credentials from gcloud."""
+
+import datetime
+import logging
+import subprocess
+from typing import Iterator, List, Optional
+
+from google.oauth2 import credentials as google_oauth2
+
+import bq_auth_flags
+import bq_flags
+import bq_utils
+from auth import utils as bq_auth_utils
+from gcloud_wrapper import gcloud_runner
+from utils import bq_error
+from utils import bq_gcloud_utils
+
+ERROR_TEXT_PRODUCED_IF_GCLOUD_NOT_FOUND = "No such file or directory: 'gcloud'"
+
+_GDRIVE_SCOPE = 'https://www.googleapis.com/auth/drive'
+_GCP_SCOPE = 'https://www.googleapis.com/auth/cloud-platform'
+
+
+def LoadCredential() -> google_oauth2.Credentials:
+  """Loads credentials by calling gcloud commands."""
+  gcloud_config = bq_gcloud_utils.load_config()
+  account = gcloud_config.get('core', {}).get('account', '')
+  logging.info('Loading auth credentials from gcloud for account: %s', account)
+
+  is_service_account = bq_utils.IsServiceAccount(account)
+  access_token = _GetAccessTokenAndPrintOutput(is_service_account)
+  # Service accounts use the refresh_handler instead of the token for refresh.
+  refresh_token = (
+      None if is_service_account else _GetRefreshTokenAndPrintOutput()
+  )
+  refresh_handler = (
+      _ServiceAccountRefreshHandler if is_service_account else None
+  )
+  fallback_quota_project_id = _GetFallbackQuotaProjectId(
+      is_service_account=is_service_account,
+      has_refresh_token=refresh_token is not None,
+  )
+
+  return google_oauth2.Credentials(
+      account=account,
+      token=access_token,
+      refresh_token=refresh_token,
+      refresh_handler=refresh_handler,
+      client_id=bq_auth_utils.get_client_id(),
+      client_secret=bq_auth_utils.get_client_secret(),
+      token_uri=bq_auth_utils.get_token_uri(),
+      quota_project_id=bq_utils.GetResolvedQuotaProjectID(
+          bq_auth_flags.QUOTA_PROJECT_ID.value, fallback_quota_project_id
+      ),
+  )
+
+
+def _GetScopes() -> List[str]:
+  scopes = []
+  if bq_flags.ENABLE_GDRIVE.value:
+    drive_scope = _GDRIVE_SCOPE
+    scopes.extend([drive_scope, _GCP_SCOPE])
+  return scopes
+
+
+def _GetAccessTokenAndPrintOutput(
+    is_service_account: bool, scopes: Optional[List[str]] = None
+) -> Optional[str]:
+  scopes = _GetScopes() if scopes is None else scopes
+  if is_service_account and scopes:
+    return _GetTokenFromGcloudAndPrintOtherOutput(
+        ['auth', 'print-access-token', '--scopes', ','.join(scopes)]
+    )
+  return _GetTokenFromGcloudAndPrintOtherOutput(['auth', 'print-access-token'])
+
+
+def _GetRefreshTokenAndPrintOutput() -> Optional[str]:
+  return _GetTokenFromGcloudAndPrintOtherOutput(['auth', 'print-refresh-token'])
+
+
+def _GetTokenFromGcloudAndPrintOtherOutput(
+    cmd: List[str],
+    stderr: Optional[int] = subprocess.STDOUT,
+) -> Optional[str]:
+  """Returns a token or prints other messages from the given gcloud command."""
+  try:
+    token = None
+    for output in _RunGcloudCommand(cmd, stderr):
+      if output and ' ' not in output:
+        # Token is a non-empty string of non-space characters.
+        token = output
+        break
+      else:
+        print(output)
+    return token
+  except bq_error.BigqueryError as e:
+    single_line_error_msg = str(e).replace('\n', '')
+    if 'security key' in single_line_error_msg:
+      raise bq_error.BigqueryError(
+          'Access token has expired. Did you touch the security key within the'
+          ' timeout window?\n'
+          + _GetReauthMessage()
+      )
+    elif 'Refresh token has expired' in single_line_error_msg:
+      raise bq_error.BigqueryError(
+          'Refresh token has expired. ' + _GetReauthMessage()
+      )
+    elif 'do not support refresh tokens' in single_line_error_msg:
+      # It's expected that certain credential types don't support refresh token.
+      return None
+    else:
+      raise bq_error.BigqueryError(
+          'Error retrieving auth credentials from gcloud: %s'
+          % _UpdateReauthMessage(str(e))
+      )
+  except Exception as e:  # pylint: disable=broad-exception-caught
+    single_line_error_msg = str(e).replace('\n', '')
+    if ERROR_TEXT_PRODUCED_IF_GCLOUD_NOT_FOUND in single_line_error_msg:
+      raise bq_error.BigqueryError(
+          "'gcloud' not found but is required for authentication. To install,"
+          ' follow these instructions:'
+          ' https://cloud.google.com/sdk/docs/install'
+      )
+    raise bq_error.BigqueryError(
+        'Error retrieving auth credentials from gcloud: %s' % str(e)
+    )
+
+
+def _RunGcloudCommand(
+    cmd: List[str], stderr: Optional[int] = subprocess.STDOUT
+) -> Iterator[str]:
+  """Runs the given gcloud command, yields the output, and returns the final status code."""
+  proc = gcloud_runner.run_gcloud_command(cmd, stderr=stderr)
+  error_msgs = []
+  if proc.stdout:
+    for stdout_line in iter(proc.stdout.readline, ''):
+      line = str(stdout_line).strip()
+      if line.startswith('ERROR:') or error_msgs:
+        error_msgs.append(line)
+      else:
+        yield line
+    proc.stdout.close()
+  return_code = proc.wait()
+  if return_code:
+    raise bq_error.BigqueryError('\n'.join(error_msgs))
+
+
+def _GetReauthMessage() -> str:
+  gcloud_command = '$ gcloud auth login' + (
+      ' --enable-gdrive-access' if bq_flags.ENABLE_GDRIVE.value else ''
+  )
+  return 'To re-authenticate, run:\n\n%s' % gcloud_command
+
+
+def _UpdateReauthMessage(message: str) -> str:
+  if '$ gcloud auth login' not in message or not bq_flags.ENABLE_GDRIVE.value:
+    return message
+  return message.replace(
+      '$ gcloud auth login',
+      '$ gcloud auth login --enable-gdrive-access',
+  )
+
+
+def _GetFallbackQuotaProjectId(
+    is_service_account: bool, has_refresh_token: bool
+) -> Optional[str]:
+  # When the credential type is not a service account - determined by the
+  # account name or whether we can get a non-empty refresh token - set a
+  # fallback quota project ID to be the resource project ID. When the credential
+  # type is a service account, don't set any fallback quota project ID.
+  if is_service_account:
+    return None
+  if not has_refresh_token:
+    return None
+  return bq_flags.PROJECT_ID.value
+
+
+def _ServiceAccountRefreshHandler(request, scopes):
+  """Refreshes the access token for a service account."""
+  del request  # Unused.
+  access_token = _GetAccessTokenAndPrintOutput(
+      is_service_account=True, scopes=scopes
+  )
+  # According to
+  # https://cloud.google.com/docs/authentication/token-types#at-lifetime
+  # and https://cloud.google.com/sdk/gcloud/reference/auth/print-access-token,
+  # the access token lifetime from gcloud auth print-access-token is 1 hour,
+  # but set token expiry to 55 minutes from now to be safe.
+  expiry = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(
+      minutes=55
+  )
+  expiry = expiry.replace(tzinfo=None)
+  return access_token, expiry

gcloud auth application-default login/google-cloud-sdk/platform/bq/auth/main_credential_loader.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python
+"""Utilities to create Google Auth credentials."""
+
+import logging
+from typing import Union
+
+from absl import app
+from google.auth.compute_engine import credentials as compute_engine
+from google.oauth2 import credentials as google_oauth2
+from google.oauth2 import service_account
+
+import bq_auth_flags
+import bq_flags
+import bq_utils
+from auth import gcloud_credential_loader
+from utils import bq_error
+
+
+GoogleAuthCredentialsUnionType = Union[
+    google_oauth2.Credentials,
+    service_account.Credentials,
+    compute_engine.Credentials,
+]
+
+
+def GetCredentialsFromFlags() -> GoogleAuthCredentialsUnionType:
+  """Returns credentials based on BQ CLI auth flags.
+
+  Returns: An OAuth2, compute engine, or service account credentials objects
+  based on BQ CLI auth flag values.
+
+  Raises:
+  app.UsageError, invalid flag values.
+  bq_error.BigqueryError, error getting credentials.
+  """
+  if bq_auth_flags.APPLICATION_DEFAULT_CREDENTIAL_FILE.value:
+    raise app.UsageError(
+        'The --application_default_credential_file flag is being deprecated.'
+        ' For now, this flag can still be used by forcing the legacy'
+        ' authentication library with --nouse_google_auth.'
+    )
+  if (
+      bq_auth_flags.SERVICE_ACCOUNT_PRIVATE_KEY_PASSWORD.default
+      != bq_auth_flags.SERVICE_ACCOUNT_PRIVATE_KEY_PASSWORD.value
+  ):
+    raise app.UsageError(bq_error.P12_DEPRECATION_MESSAGE)
+
+  if bq_auth_flags.OAUTH_ACCESS_TOKEN.value:
+    logging.info('Loading auth credentials from --oauth_access_token')
+    return google_oauth2.Credentials(
+        token=bq_auth_flags.OAUTH_ACCESS_TOKEN.value,
+        quota_project_id=bq_utils.GetResolvedQuotaProjectID(
+            bq_auth_flags.QUOTA_PROJECT_ID.value, bq_flags.PROJECT_ID.value
+        ),
+    )
+  else:
+    logging.info('No `oauth_access_token`, load credentials elsewhere')
+
+  if bq_auth_flags.USE_GCE_SERVICE_ACCOUNT.value:
+    logging.info('Loading auth credentials with --use_gce_service_account')
+    return compute_engine.Credentials(
+        quota_project_id=bq_utils.GetResolvedQuotaProjectID(
+            bq_auth_flags.QUOTA_PROJECT_ID.value, bq_flags.PROJECT_ID.value
+        ),
+    )
+  else:
+    logging.info('No `use_gce_service_account`, load credentials elsewhere')
+
+  if bq_auth_flags.SERVICE_ACCOUNT.value:
+    raise app.UsageError(
+        'The flag --service_account is not supported. '
+        'To use a service account please follow'
+        ' https://cloud.google.com/docs/authentication/'
+        'use-service-account-impersonation#gcloud-config.'
+    )
+
+
+  return gcloud_credential_loader.LoadCredential()
+
+

gcloud auth application-default login/google-cloud-sdk/platform/bq/auth/utils.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python
+"""BQ CLI helper functions for authentication."""
+
+import os
+
+import bq_flags
+import bq_utils
+
+
+_DEFAULT_TOKEN_HOST = 'https://oauth2.googleapis.com/token'
+_DEFAULT_MTLS_TOKEN_HOST = 'https://oauth2.mtls.googleapis.com/token'
+
+
+_CLOUD_CLI_CLIENT_ID = '32555940559.apps.googleusercontent.com'
+_CLOUD_CLI_CLIENT_SECRET = 'ZmssLNjJy2998hD4CTg2ejr2'
+_CLOUD_CLI_CLIENT_USER_AGENT = 'google-cloud-sdk' + os.environ.get(
+    'CLOUDSDK_VERSION', bq_utils.VERSION_NUMBER
+)
+
+_BQ_CLI_CLIENT_ID = '977385342095.apps.googleusercontent.com'
+_BQ_CLI_CLIENT_SECRET = 'wbER7576mc_1YOII0dGk7jEE'
+_BQ_CLI_CLIENT_USER_AGENT = 'bq/' + bq_utils.VERSION_NUMBER
+
+
+def get_client_id() -> str:
+  if os.environ.get('CLOUDSDK_WRAPPER') == '1':
+    return _CLOUD_CLI_CLIENT_ID
+  else:
+    return _BQ_CLI_CLIENT_ID
+
+
+def get_client_secret() -> str:
+  if os.environ.get('CLOUDSDK_WRAPPER') == '1':
+    return _CLOUD_CLI_CLIENT_SECRET
+  else:
+    return _BQ_CLI_CLIENT_SECRET
+
+
+def get_client_user_agent() -> str:
+  if os.environ.get('CLOUDSDK_WRAPPER') == '1':
+    return _CLOUD_CLI_CLIENT_USER_AGENT
+  else:
+    return _BQ_CLI_CLIENT_USER_AGENT
+
+
+def get_token_uri() -> str:
+  # TODO: b/394417874 - Read user's gcloud properties too.
+  if bq_flags.MTLS.value:
+    return _DEFAULT_MTLS_TOKEN_HOST
+  return _DEFAULT_TOKEN_HOST