import asyncio import calendar import json import logging import os import shlex import subprocess import time from typing import Callable, Optional from alibabacloud_credentials.provider.refreshable import Credentials, RefreshResult, RefreshCachedSupplier from alibabacloud_credentials_api import ICredentialsProvider from alibabacloud_credentials.exceptions import CredentialException log = logging.getLogger('credentials') ExternalCredentialUpdateCallback = Callable[[str, str, str, int], None] ExternalCredentialUpdateCallbackAsync = Callable[[str, str, str, int], None] def _parse_expiration(expiration: str) -> int: if not expiration: return 0 time_array = time.strptime(expiration, '%Y-%m-%dT%H:%M:%SZ') return calendar.timegm(time_array) def _get_stale_time(expiration: int) -> int: if expiration <= 0: return int(time.mktime(time.localtime())) return expiration - 180 class ExternalCredentialsProvider(ICredentialsProvider): DEFAULT_TIMEOUT = 60 def __init__(self, *, process_command: str = None, timeout: int = None, credential_update_callback: Optional[ExternalCredentialUpdateCallback] = None, credential_update_callback_async: Optional[ExternalCredentialUpdateCallbackAsync] = None): if not process_command: raise ValueError('process_command is empty') self._process_command = process_command self._timeout = timeout if timeout and timeout > 0 else ExternalCredentialsProvider.DEFAULT_TIMEOUT self._credential_update_callback = credential_update_callback self._credential_update_callback_async = credential_update_callback_async self._credentials_cache = RefreshCachedSupplier( refresh_callable=self._refresh_credentials, refresh_callable_async=self._refresh_credentials_async, ) def get_credentials(self) -> Credentials: return self._credentials_cache._sync_call() async def get_credentials_async(self) -> Credentials: return await self._credentials_cache._async_call() def _refresh_credentials(self) -> RefreshResult[Credentials]: if not self._process_command.strip(): raise CredentialException('process_command is empty') try: command = self._process_command if os.name == 'nt' else shlex.split(self._process_command) completed = subprocess.run( command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=self._timeout, check=False, text=True, shell=os.name == 'nt', ) except subprocess.TimeoutExpired: raise CredentialException(f'command process timed out after {self._timeout * 1000} milliseconds') except Exception as e: raise CredentialException(f'failed to execute external command: {e}') if completed.returncode != 0: raise CredentialException( f'failed to execute external command: exit status {completed.returncode}\nstderr: {completed.stderr}') return self._parse_and_build_credentials(completed.stdout, async_callback=False) async def _refresh_credentials_async(self) -> RefreshResult[Credentials]: if not self._process_command.strip(): raise CredentialException('process_command is empty') try: if os.name == 'nt': process = await asyncio.create_subprocess_shell( self._process_command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) else: process = await asyncio.create_subprocess_exec( *shlex.split(self._process_command), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, ) stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=self._timeout) except asyncio.TimeoutError: if 'process' in locals(): process.kill() await process.wait() raise CredentialException(f'command process timed out after {self._timeout * 1000} milliseconds') except Exception as e: raise CredentialException(f'failed to execute external command: {e}') if process.returncode != 0: raise CredentialException( f'failed to execute external command: exit status {process.returncode}\nstderr: {stderr.decode("utf-8")}') return await self._parse_and_build_credentials_async(stdout.decode('utf-8')) def _parse_and_build_credentials(self, output: str, async_callback: bool) -> RefreshResult[Credentials]: try: data = json.loads(output) except Exception as e: raise CredentialException(f'failed to parse external command output: {e}') access_key_id = data.get('access_key_id') access_key_secret = data.get('access_key_secret') security_token = data.get('sts_token') if not access_key_id or not access_key_secret: raise CredentialException('invalid credential response: access_key_id or access_key_secret is empty') if data.get('mode') == 'StsToken' and not security_token: raise CredentialException('invalid StsToken credential response: sts_token is empty') expiration = _parse_expiration(data.get('expiration')) credentials = Credentials( access_key_id=access_key_id, access_key_secret=access_key_secret, security_token=security_token, expiration=expiration, provider_name=self.get_provider_name(), ) if not async_callback and self._credential_update_callback: try: self._credential_update_callback(access_key_id, access_key_secret, security_token, expiration) except Exception as e: log.warning(f'failed to update external credentials in config file: {e}') return RefreshResult(value=credentials, stale_time=_get_stale_time(expiration)) async def _parse_and_build_credentials_async(self, output: str) -> RefreshResult[Credentials]: result = self._parse_and_build_credentials(output, async_callback=True) credentials = result.value() if self._credential_update_callback_async: try: await self._credential_update_callback_async( credentials.get_access_key_id(), credentials.get_access_key_secret(), credentials.get_security_token(), credentials.get_expiration() or 0, ) except Exception as e: log.warning(f'failed to update external credentials in config file: {e}') return result def get_provider_name(self) -> str: return 'external'