feat: support OAuth credentials provider

Author: yndu13

alibabacloud_credentials/provider/__init__.py

@@ -10,6 +10,7 @@
 from .profile import ProfileCredentialsProvider
 from .default import DefaultCredentialsProvider
 from .cloud_sso import CloudSSOCredentialsProvider
+from .oauth import OAuthCredentialsProvider
 
 __all__ = [
     'StaticAKCredentialsProvider',
@@ -23,5 +24,6 @@
     'CLIProfileCredentialsProvider',
     'ProfileCredentialsProvider',
     'DefaultCredentialsProvider',
-    'CloudSSOCredentialsProvider'
+    'CloudSSOCredentialsProvider',
+    'OAuthCredentialsProvider'
 ]

alibabacloud_credentials/provider/cli_profile.py

@@ -1,5 +1,7 @@
 import os
 import json
+import threading
+import fcntl
 from typing import Any, Dict
 
 import aiofiles
@@ -10,6 +12,7 @@
 from .oidc import OIDCRoleArnCredentialsProvider
 from .static_sts import StaticSTSCredentialsProvider
 from .cloud_sso import CloudSSOCredentialsProvider
+from .oauth import OAuthCredentialsProvider, OAuthTokenUpdateCallback, OAuthTokenUpdateCallbackAsync
 from .refreshable import Credentials
 from alibabacloud_credentials_api import ICredentialsProvider
 from alibabacloud_credentials.utils import auth_constant as ac
@@ -32,10 +35,13 @@ def _load_config(file_path: str) -> Any:
 class CLIProfileCredentialsProvider(ICredentialsProvider):
 
     def __init__(self, *,
-                 profile_name: str = None):
-        self._profile_file = os.path.join(ac.HOME, ".aliyun/config.json")
+                 profile_name: str = None,
+                 profile_file: str = None):
+        self._profile_file = profile_file or os.path.join(ac.HOME, ".aliyun/config.json")
         self._profile_name = profile_name or au.environment_profile_name
         self.__innerProvider = None
+        # 文件锁，用于并发安全
+        self._file_lock = threading.RLock()
 
     def _should_reload_credentials_provider(self) -> bool:
         if self.__innerProvider is None:
@@ -175,10 +181,225 @@ def _get_credentials_provider(self, config: Dict, profile_name: str) -> ICredent
                         access_token=profile.get('access_token'),
                         access_token_expire=profile.get('cloud_sso_access_token_expire'),
                     )
+                elif mode == "OAuth":
+                    # 获取 OAuth 配置
+                    site_type = profile.get('oauth_site_type', 'CN')
+                    oauth_base_url_map = {
+                        'CN': 'https://oauth.aliyun.com',
+                        'INTL': 'https://oauth.alibabacloud.com'
+                    }
+                    sign_in_url = oauth_base_url_map.get(site_type.upper())
+                    if not sign_in_url:
+                        raise CredentialException('Invalid OAuth site type, support CN or INTL')
+
+                    oauth_client_map = {
+                        'CN': '4038181954557748008',
+                        'INTL': '4103531455503354461'
+                    }
+                    client_id = oauth_client_map.get(site_type.upper())
+                    if not client_id:
+                        raise CredentialException('Invalid OAuth site type, support CN or INTL')
+
+                    return OAuthCredentialsProvider(
+                        client_id=client_id,
+                        sign_in_url=sign_in_url,
+                        access_token=profile.get('oauth_access_token'),
+                        access_token_expire=profile.get('oauth_access_token_expire'),
+                        refresh_token=profile.get('oauth_refresh_token'),
+                        token_update_callback=self._get_oauth_token_update_callback(),
+                        token_update_callback_async=self._get_oauth_token_update_callback_async(),
+                    )
                 else:
                     raise CredentialException(f"unsupported profile mode '{mode}' form cli credentials file.")
 
         raise CredentialException(f"unable to get profile with '{profile_name}' form cli credentials file.")
 
     def get_provider_name(self) -> str:
         return 'cli_profile'
+
+    def _update_oauth_tokens(self, refresh_token: str, access_token: str, access_key: str, secret: str,
+                             security_token: str, access_token_expire: int, sts_expire: int) -> None:
+        """更新 OAuth 令牌并写回配置文件"""
+        with self._file_lock:
+            try:
+                # 读取现有配置
+                config = _load_config(self._profile_file)
+
+                # 找到当前 profile 并更新 OAuth 令牌
+                profile_name = self._profile_name
+                if not profile_name:
+                    profile_name = config.get('current')
+                profiles = config.get('profiles', [])
+                profile_tag = False
+                for profile in profiles:
+                    if profile.get('name') == profile_name:
+                        profile_tag = True
+                        # 更新 OAuth 令牌
+                        profile['oauth_refresh_token'] = refresh_token
+                        profile['oauth_access_token'] = access_token
+                        profile['oauth_access_token_expire'] = access_token_expire
+                        # 更新 STS 凭据
+                        profile['access_key_id'] = access_key
+                        profile['access_key_secret'] = secret
+                        profile['sts_token'] = security_token
+                        profile['sts_expiration'] = sts_expire
+                        break
+
+                # 写回配置文件
+                if not profile_tag:
+                    raise CredentialException(f"unable to get profile with '{profile_name}' form cli credentials file.")
+
+                self._write_configuration_to_file(self._profile_file, config)
+
+            except Exception as e:
+                raise CredentialException(f"failed to update OAuth tokens in config file: {e}")
+
+    def _write_configuration_to_file(self, config_path: str, config: Dict) -> None:
+        """将配置写入文件，使用原子写入确保数据完整性"""
+        temp_file = config_path + '.tmp'
+        try:
+            # 序列化配置
+            data = json.dumps(config, indent=4, ensure_ascii=False)
+
+            # 写入临时文件
+            with open(temp_file, 'w', encoding='utf-8') as f:
+                f.write(data)
+
+            # 原子性重命名，确保文件完整性
+            os.rename(temp_file, config_path)
+
+        except Exception as e:
+            # 清理临时文件
+            if os.path.exists(temp_file):
+                os.remove(temp_file)
+            raise e
+
+    def _write_configuration_to_file_with_lock(self, config_path: str, config: Dict) -> None:
+        """使用操作系统级别的文件锁写入配置文件"""
+        try:
+            # 打开文件用于锁定
+            with open(config_path, 'r+') as f:
+                # 获取独占锁（阻塞其他进程）
+                fcntl.flock(f.fileno(), fcntl.LOCK_EX)
+
+                try:
+                    # 序列化配置
+                    data = json.dumps(config, indent=4, ensure_ascii=False)
+
+                    # 创建临时文件
+                    temp_file = config_path + '.tmp'
+                    with open(temp_file, 'w', encoding='utf-8') as temp_f:
+                        temp_f.write(data)
+
+                    # 原子性重命名
+                    os.rename(temp_file, config_path)
+
+                finally:
+                    # 释放锁
+                    fcntl.flock(f.fileno(), fcntl.LOCK_UN)
+
+        except Exception as e:
+            # 清理临时文件
+            temp_file = config_path + '.tmp'
+            if os.path.exists(temp_file):
+                os.remove(temp_file)
+            raise e
+
+    def _get_oauth_token_update_callback(self) -> OAuthTokenUpdateCallback:
+        """获取 OAuth 令牌更新回调函数"""
+        return lambda refresh_token, access_token, access_key, secret, security_token, access_token_expire, sts_expire: self._update_oauth_tokens(
+            refresh_token, access_token, access_key, secret, security_token, access_token_expire, sts_expire
+        )
+
+    async def _write_configuration_to_file_async(self, config_path: str, config: Dict) -> None:
+        """异步将配置写入文件，使用原子写入确保数据完整性"""
+        temp_file = config_path + '.tmp'
+        try:
+            # 序列化配置
+            data = json.dumps(config, indent=4, ensure_ascii=False)
+
+            # 异步写入临时文件
+            async with aiofiles.open(temp_file, 'w', encoding='utf-8') as f:
+                await f.write(data)
+
+            # 原子性重命名
+            os.rename(temp_file, config_path)
+
+        except Exception as e:
+            if os.path.exists(temp_file):
+                os.remove(temp_file)
+            raise e
+
+    async def _write_configuration_to_file_with_lock_async(self, config_path: str, config: Dict) -> None:
+        """异步使用操作系统级别的文件锁写入配置文件"""
+        try:
+            # 打开文件用于锁定
+            with open(config_path, 'r+') as f:
+                # 获取独占锁（阻塞其他进程）
+                fcntl.flock(f.fileno(), fcntl.LOCK_EX)
+
+                try:
+                    # 序列化配置
+                    data = json.dumps(config, indent=4, ensure_ascii=False)
+
+                    # 创建临时文件
+                    temp_file = config_path + '.tmp'
+                    async with aiofiles.open(temp_file, 'w', encoding='utf-8') as temp_f:
+                        await temp_f.write(data)
+
+                    # 原子性重命名
+                    os.rename(temp_file, config_path)
+
+                finally:
+                    # 释放锁
+                    fcntl.flock(f.fileno(), fcntl.LOCK_UN)
+
+        except Exception as e:
+            # 清理临时文件
+            temp_file = config_path + '.tmp'
+            if os.path.exists(temp_file):
+                os.remove(temp_file)
+            raise e
+
+    async def _update_oauth_tokens_async(self, refresh_token: str, access_token: str, access_key: str, secret: str,
+                                         security_token: str, access_token_expire: int, sts_expire: int) -> None:
+        """异步更新 OAuth 令牌并写回配置文件"""
+        try:
+            with self._file_lock:
+                cfg_path = self._profile_file
+                conf = await _load_config_async(cfg_path)
+
+                # 找到当前 profile 并更新 OAuth 令牌
+                profile_name = self._profile_name
+                if not profile_name:
+                    profile_name = conf.get('current')
+                profiles = conf.get('profiles', [])
+                profile_tag = False
+                for profile in profiles:
+                    if profile.get('name') == profile_name:
+                        profile_tag = True
+                        # 更新 OAuth 相关字段
+                        profile['oauth_refresh_token'] = refresh_token
+                        profile['oauth_access_token'] = access_token
+                        profile['oauth_access_token_expire'] = access_token_expire
+                        # 更新 STS 凭据
+                        profile['access_key_id'] = access_key
+                        profile['access_key_secret'] = secret
+                        profile['sts_token'] = security_token
+                        profile['sts_expiration'] = sts_expire
+                        break
+
+                if not profile_tag:
+                    raise CredentialException(f"Profile '{profile_name}' not found in config file")
+
+                # 异步写回配置文件
+                await self._write_configuration_to_file_with_lock_async(cfg_path, conf)
+
+        except Exception as e:
+            raise CredentialException(f"failed to update OAuth tokens in config file: {e}")
+
+    def _get_oauth_token_update_callback_async(self) -> OAuthTokenUpdateCallbackAsync:
+        """获取异步 OAuth 令牌更新回调函数"""
+        return lambda refresh_token, access_token, access_key, secret, security_token, access_token_expire, sts_expire: self._update_oauth_tokens_async(
+            refresh_token, access_token, access_key, secret, security_token, access_token_expire, sts_expire
+        )