import calendar import json import logging import time from urllib.parse import urlparse, urlencode from typing import Callable, Optional from alibabacloud_credentials.provider.refreshable import Credentials, RefreshResult, RefreshCachedSupplier from alibabacloud_credentials.http import HttpOptions from Tea.core import TeaCore from alibabacloud_credentials_api import ICredentialsProvider from alibabacloud_credentials.utils import parameter_helper as ph from alibabacloud_credentials.exceptions import CredentialException log = logging.getLogger('credentials') log.setLevel(logging.INFO) ch = logging.StreamHandler() log.addHandler(ch) # OAuth 令牌更新回调函数类型 OAuthTokenUpdateCallback = Callable[[str, str, str, str, str, int, int], None] OAuthTokenUpdateCallbackAsync = Callable[[str, str, str, str, str, int, int], None] def _get_stale_time(expiration: int) -> int: if expiration < 0: return int(time.mktime(time.localtime())) + 60 * 60 # OAuth STS credentials are persisted in the CLI profile and should be # reused for their complete lifetime. Using an earlier stale time causes # every credentials lookup near expiration to enter the refresh path even # though the persisted STS credential is still valid. return expiration class OAuthCredentialsProvider(ICredentialsProvider): DEFAULT_CONNECT_TIMEOUT = 5000 DEFAULT_READ_TIMEOUT = 10000 def __init__(self, *, client_id: str = None, sign_in_url: str = None, access_token: str = None, access_token_expire: int = 0, refresh_token: str = None, access_key_id: str = None, access_key_secret: str = None, security_token: str = None, sts_expiration: int = None, http_options: HttpOptions = None, token_update_callback: Optional[OAuthTokenUpdateCallback] = None, token_update_callback_async: Optional[OAuthTokenUpdateCallbackAsync] = None): if not client_id: raise ValueError('the ClientId is empty') if not sign_in_url: raise ValueError('the url for sign-in is empty') self._client_id = client_id self._sign_in_url = sign_in_url self._access_token = access_token self._access_token_expire = access_token_expire self._refresh_token = refresh_token # Cached STS credential previously obtained from a token exchange and # persisted alongside the OAuth tokens (e.g. in the CLI config file). # Reusing it until ``sts_expiration`` avoids exchanging on every fresh # provider build, matching aliyun-cli's OAuth resolution. self._sts_access_key_id = access_key_id self._sts_access_key_secret = access_key_secret self._sts_security_token = security_token self._sts_expiration = sts_expiration self._token_update_callback = token_update_callback self._token_update_callback_async = token_update_callback_async self._http_options = http_options if http_options is not None else HttpOptions() self._runtime_options = { 'connectTimeout': self._http_options.connect_timeout if self._http_options.connect_timeout is not None else OAuthCredentialsProvider.DEFAULT_CONNECT_TIMEOUT, 'readTimeout': self._http_options.read_timeout if self._http_options.read_timeout is not None else OAuthCredentialsProvider.DEFAULT_READ_TIMEOUT, 'httpsProxy': self._http_options.proxy } 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 _reuse_cached_sts(self) -> Optional[RefreshResult]: """Return the cached STS credential when it is still valid. Mirrors aliyun-cli's OAuth handling (config/profile.go): if a previously exchanged STS credential is present and not yet expired, reuse it instead of calling ``/v1/exchange``. This keeps a rotating STS AccessKey id stable across fresh provider builds and, crucially, avoids re-triggering the token-update callback (which rewrites the config file) on every resolve. """ if (self._sts_access_key_id and self._sts_access_key_secret and self._sts_security_token and self._sts_expiration and self._sts_expiration > int(time.mktime(time.localtime()))): credentials = Credentials( access_key_id=self._sts_access_key_id, access_key_secret=self._sts_access_key_secret, security_token=self._sts_security_token, expiration=self._sts_expiration, provider_name=self.get_provider_name(), ) return RefreshResult(value=credentials, stale_time=_get_stale_time(self._sts_expiration)) return None def _store_exchanged_sts(self, credentials: Credentials, expiration: int) -> None: """Remember the freshly exchanged STS so later resolves can reuse it.""" self._sts_access_key_id = credentials.get_access_key_id() self._sts_access_key_secret = credentials.get_access_key_secret() self._sts_security_token = credentials.get_security_token() self._sts_expiration = expiration def _try_refresh_oauth_token(self) -> None: current_time = int(time.mktime(time.localtime())) # 构建刷新令牌请求 r = urlparse(self._sign_in_url) tea_request = ph.get_new_request() tea_request.headers['host'] = r.hostname tea_request.port = r.port tea_request.protocol = r.scheme tea_request.method = 'POST' tea_request.pathname = '/v1/token' # 设置请求体 body_data = { 'grant_type': 'refresh_token', 'refresh_token': self._refresh_token, 'client_id': self._client_id, 'Timestamp': ph.get_iso_8061_date() } tea_request.body = urlencode(body_data) tea_request.headers['Content-Type'] = 'application/x-www-form-urlencoded' response = TeaCore.do_action(tea_request, self._runtime_options) if response.status_code != 200: raise CredentialException(f"failed to refresh OAuth token, status code: {response.status_code}, response: {response.body.decode('utf-8')}") # 解析响应 dic = json.loads(response.body.decode('utf-8')) if 'access_token' not in dic or 'refresh_token' not in dic: raise CredentialException(f"failed to refresh OAuth token: {response.body.decode('utf-8')}") # 更新令牌 new_access_token = dic.get('access_token') new_refresh_token = dic.get('refresh_token') expires_in = dic.get('expires_in', 3600) new_access_token_expire = current_time + expires_in self._access_token = new_access_token self._refresh_token = new_refresh_token self._access_token_expire = new_access_token_expire async def _try_refresh_oauth_token_async(self) -> None: current_time = int(time.mktime(time.localtime())) # 构建刷新令牌请求 r = urlparse(self._sign_in_url) tea_request = ph.get_new_request() tea_request.headers['host'] = r.hostname tea_request.port = r.port tea_request.protocol = r.scheme tea_request.method = 'POST' tea_request.pathname = '/v1/token' # 设置请求体 body_data = { 'grant_type': 'refresh_token', 'refresh_token': self._refresh_token, 'client_id': self._client_id, 'Timestamp': ph.get_iso_8061_date() } tea_request.body = urlencode(body_data) tea_request.headers['Content-Type'] = 'application/x-www-form-urlencoded' response = await TeaCore.async_do_action(tea_request, self._runtime_options) if response.status_code != 200: raise CredentialException(f"failed to refresh OAuth token, status code: {response.status_code}, response: {response.body.decode('utf-8')}") # 解析响应 dic = json.loads(response.body.decode('utf-8')) if 'access_token' not in dic or 'refresh_token' not in dic: raise CredentialException(f"failed to refresh OAuth token: {response.body.decode('utf-8')}") # 更新令牌 new_access_token = dic.get('access_token') new_refresh_token = dic.get('refresh_token') expires_in = dic.get('expires_in', 3600) new_access_token_expire = current_time + expires_in self._access_token = new_access_token self._refresh_token = new_refresh_token self._access_token_expire = new_access_token_expire def _refresh_credentials(self) -> RefreshResult[Credentials]: cached = self._reuse_cached_sts() if cached is not None: return cached if self._refresh_token and ( self._access_token is None or self._access_token_expire <= 0 or self._access_token_expire - int( time.mktime(time.localtime())) <= 1200): self._try_refresh_oauth_token() r = urlparse(self._sign_in_url) tea_request = ph.get_new_request() tea_request.headers['host'] = r.hostname tea_request.port = r.port tea_request.protocol = r.scheme tea_request.method = 'POST' tea_request.pathname = '/v1/exchange' tea_request.headers['Content-Type'] = 'application/json' tea_request.headers['Authorization'] = f'Bearer {self._access_token}' response = TeaCore.do_action(tea_request, self._runtime_options) if response.status_code != 200: raise CredentialException( f"error refreshing credentials from OAuth, http_code: {response.status_code}, result: {response.body.decode('utf-8')}") dic = json.loads(response.body.decode('utf-8')) if 'error' in dic: raise CredentialException( f"error retrieving credentials from OAuth result: {response.body.decode('utf-8')}") if 'AccessKeyId' not in dic or 'AccessKeySecret' not in dic or 'SecurityToken' not in dic: raise CredentialException( f"error retrieving credentials from OAuth result: {response.body.decode('utf-8')}") # 先转换为时间数组 time_array = time.strptime(dic.get('Expiration'), '%Y-%m-%dT%H:%M:%SZ') # 转换为时间戳 expiration = calendar.timegm(time_array) credentials = Credentials( access_key_id=dic.get('AccessKeyId'), access_key_secret=dic.get('AccessKeySecret'), security_token=dic.get('SecurityToken'), expiration=expiration, provider_name=self.get_provider_name() ) self._store_exchanged_sts(credentials, expiration) # 调用令牌更新回调函数 if self._token_update_callback: try: self._token_update_callback( self._refresh_token, self._access_token, credentials.get_access_key_id(), credentials.get_access_key_secret(), credentials.get_security_token(), self._access_token_expire, expiration ) except Exception as e: log.warning(f'failed to update OAuth tokens in config file: {e}') return RefreshResult(value=credentials, stale_time=_get_stale_time(expiration)) async def _refresh_credentials_async(self) -> RefreshResult[Credentials]: cached = self._reuse_cached_sts() if cached is not None: return cached if self._refresh_token and ( self._access_token is None or self._access_token_expire <= 0 or self._access_token_expire - int( time.mktime(time.localtime())) <= 1200): await self._try_refresh_oauth_token_async() r = urlparse(self._sign_in_url) tea_request = ph.get_new_request() tea_request.headers['host'] = r.hostname tea_request.port = r.port tea_request.protocol = r.scheme tea_request.method = 'POST' tea_request.pathname = '/v1/exchange' tea_request.headers['Content-Type'] = 'application/json' tea_request.headers['Authorization'] = f'Bearer {self._access_token}' response = await TeaCore.async_do_action(tea_request, self._runtime_options) if response.status_code != 200: raise CredentialException( f"error refreshing credentials from OAuth, http_code: {response.status_code}, result: {response.body.decode('utf-8')}") dic = json.loads(response.body.decode('utf-8')) if 'error' in dic: raise CredentialException( f"error retrieving credentials from OAuth result: {response.body.decode('utf-8')}") if 'AccessKeyId' not in dic or 'AccessKeySecret' not in dic or 'SecurityToken' not in dic: raise CredentialException( f"error retrieving credentials from OAuth result: {response.body.decode('utf-8')}") # 先转换为时间数组 time_array = time.strptime(dic.get('Expiration'), '%Y-%m-%dT%H:%M:%SZ') # 转换为时间戳 expiration = calendar.timegm(time_array) credentials = Credentials( access_key_id=dic.get('AccessKeyId'), access_key_secret=dic.get('AccessKeySecret'), security_token=dic.get('SecurityToken'), expiration=expiration, provider_name=self.get_provider_name() ) self._store_exchanged_sts(credentials, expiration) if self._token_update_callback_async: try: await self._token_update_callback_async( self._refresh_token, self._access_token, credentials.get_access_key_id(), credentials.get_access_key_secret(), credentials.get_security_token(), self._access_token_expire, expiration ) except Exception as e: log.warning(f'failed to update OAuth tokens in config file: {e}') return RefreshResult(value=credentials, stale_time=_get_stale_time(expiration)) def _get_client_id(self) -> str: """获取客户端ID""" return self._client_id def get_provider_name(self) -> str: return 'oauth'