cli_profile.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701
  1. import os
  2. import json
  3. import threading
  4. import platform
  5. from typing import Any, Dict
  6. import aiofiles
  7. # 跨平台文件锁支持
  8. if platform.system() == 'Windows':
  9. # Windows平台使用msvcrt
  10. import msvcrt
  11. HAS_MSVCRT = True
  12. HAS_FCNTL = False
  13. else:
  14. # 其他平台尝试使用fcntl,如果不可用则不设文件锁
  15. HAS_MSVCRT = False
  16. try:
  17. import fcntl
  18. HAS_FCNTL = True
  19. except ImportError:
  20. HAS_FCNTL = False
  21. from .static_ak import StaticAKCredentialsProvider
  22. from .ecs_ram_role import EcsRamRoleCredentialsProvider
  23. from .ram_role_arn import RamRoleArnCredentialsProvider
  24. from .oidc import OIDCRoleArnCredentialsProvider
  25. from .static_sts import StaticSTSCredentialsProvider
  26. from .cloud_sso import CloudSSOCredentialsProvider
  27. from .oauth import OAuthCredentialsProvider, OAuthTokenUpdateCallback, OAuthTokenUpdateCallbackAsync
  28. from .external import (
  29. ExternalCredentialsProvider,
  30. ExternalCredentialUpdateCallback,
  31. ExternalCredentialUpdateCallbackAsync,
  32. )
  33. from .refreshable import Credentials
  34. from alibabacloud_credentials_api import ICredentialsProvider
  35. from alibabacloud_credentials.utils import auth_constant as ac
  36. from alibabacloud_credentials.utils import auth_util as au
  37. from alibabacloud_credentials.exceptions import CredentialException
  38. async def _load_config_async(file_path: str) -> Any:
  39. async with aiofiles.open(file_path, mode='r') as f:
  40. content = await f.read()
  41. return json.loads(content)
  42. def _load_config(file_path: str) -> Any:
  43. with open(file_path, mode='r') as f:
  44. content = f.read()
  45. return json.loads(content)
  46. class CLIProfileCredentialsProvider(ICredentialsProvider):
  47. def __init__(self, *,
  48. profile_name: str = None,
  49. profile_file: str = None,
  50. allow_config_force_rewrite: bool = False):
  51. self._profile_file = profile_file or os.path.join(ac.HOME, ".aliyun/config.json")
  52. self._profile_name = profile_name or au.environment_profile_name
  53. self._allow_config_force_rewrite = allow_config_force_rewrite
  54. self.__innerProvider = None
  55. # 文件锁,用于并发安全
  56. self._file_lock = threading.RLock()
  57. def _should_reload_credentials_provider(self) -> bool:
  58. if self.__innerProvider is None:
  59. return True
  60. return False
  61. def get_credentials(self) -> Credentials:
  62. if au.environment_cli_profile_disabled.lower() == "true":
  63. raise CredentialException('cli credentials file is disabled')
  64. if self._should_reload_credentials_provider():
  65. if not os.path.exists(self._profile_file) or not os.path.isfile(self._profile_file):
  66. raise CredentialException(f'unable to open credentials file: {self._profile_file}')
  67. try:
  68. config = _load_config(self._profile_file)
  69. except Exception as e:
  70. raise CredentialException(
  71. f'failed to parse credential form cli credentials file: {self._profile_file}')
  72. if config is None:
  73. raise CredentialException(
  74. f'failed to parse credential form cli credentials file: {self._profile_file}')
  75. profile_name = self._profile_name
  76. if self._profile_name is None or self._profile_name == '':
  77. profile_name = config.get('current')
  78. self.__innerProvider = self._get_credentials_provider(config, profile_name)
  79. cre = self.__innerProvider.get_credentials()
  80. credentials = Credentials(
  81. access_key_id=cre.get_access_key_id(),
  82. access_key_secret=cre.get_access_key_secret(),
  83. security_token=cre.get_security_token(),
  84. provider_name=f'{self.get_provider_name()}/{cre.get_provider_name()}'
  85. )
  86. return credentials
  87. async def get_credentials_async(self) -> Credentials:
  88. if au.environment_cli_profile_disabled.lower() == "true":
  89. raise CredentialException('cli credentials file is disabled')
  90. if self._should_reload_credentials_provider():
  91. if not os.path.exists(self._profile_file) or not os.path.isfile(self._profile_file):
  92. raise CredentialException(f'unable to open credentials file: {self._profile_file}')
  93. try:
  94. config = await _load_config_async(self._profile_file)
  95. except Exception as e:
  96. raise CredentialException(
  97. f'failed to parse credential form cli credentials file: {self._profile_file}')
  98. if config is None:
  99. raise CredentialException(
  100. f'failed to parse credential form cli credentials file: {self._profile_file}')
  101. profile_name = self._profile_name
  102. if self._profile_name is None or self._profile_name == '':
  103. profile_name = config.get('current')
  104. self.__innerProvider = self._get_credentials_provider(config, profile_name)
  105. cre = await self.__innerProvider.get_credentials_async()
  106. credentials = Credentials(
  107. access_key_id=cre.get_access_key_id(),
  108. access_key_secret=cre.get_access_key_secret(),
  109. security_token=cre.get_security_token(),
  110. provider_name=f'{self.get_provider_name()}/{cre.get_provider_name()}'
  111. )
  112. return credentials
  113. def _get_credentials_provider(self, config: Dict, profile_name: str) -> ICredentialsProvider:
  114. if profile_name is None or profile_name == '':
  115. raise CredentialException('invalid profile name')
  116. profiles = config.get('profiles', [])
  117. if not profiles:
  118. raise CredentialException(f"unable to get profile with '{profile_name}' form cli credentials file.")
  119. for profile in profiles:
  120. if profile.get('name') is not None and profile['name'] == profile_name:
  121. mode = profile.get('mode')
  122. if mode == "AK":
  123. return StaticAKCredentialsProvider(
  124. access_key_id=profile.get('access_key_id'),
  125. access_key_secret=profile.get('access_key_secret')
  126. )
  127. elif mode == "StsToken":
  128. return StaticSTSCredentialsProvider(
  129. access_key_id=profile.get('access_key_id'),
  130. access_key_secret=profile.get('access_key_secret'),
  131. security_token=profile.get('sts_token')
  132. )
  133. elif mode == "RamRoleArn":
  134. pre_provider = StaticAKCredentialsProvider(
  135. access_key_id=profile.get('access_key_id'),
  136. access_key_secret=profile.get('access_key_secret')
  137. )
  138. return RamRoleArnCredentialsProvider(
  139. credentials_provider=pre_provider,
  140. role_arn=profile.get('ram_role_arn'),
  141. role_session_name=profile.get('ram_session_name'),
  142. duration_seconds=profile.get('expired_seconds'),
  143. policy=profile.get('policy'),
  144. external_id=profile.get('external_id'),
  145. sts_region_id=profile.get('sts_region'),
  146. enable_vpc=profile.get('enable_vpc'),
  147. )
  148. elif mode == "EcsRamRole":
  149. return EcsRamRoleCredentialsProvider(
  150. role_name=profile.get('ram_role_name')
  151. )
  152. elif mode == "OIDC":
  153. return OIDCRoleArnCredentialsProvider(
  154. role_arn=profile.get('ram_role_arn'),
  155. oidc_provider_arn=profile.get('oidc_provider_arn'),
  156. oidc_token_file_path=profile.get('oidc_token_file'),
  157. role_session_name=profile.get('role_session_name'),
  158. duration_seconds=profile.get('expired_seconds'),
  159. policy=profile.get('policy'),
  160. sts_region_id=profile.get('sts_region'),
  161. enable_vpc=profile.get('enable_vpc'),
  162. )
  163. elif mode == "ChainableRamRoleArn":
  164. previous_provider = self._get_credentials_provider(config, profile.get('source_profile'))
  165. return RamRoleArnCredentialsProvider(
  166. credentials_provider=previous_provider,
  167. role_arn=profile.get('ram_role_arn'),
  168. role_session_name=profile.get('ram_session_name'),
  169. duration_seconds=profile.get('expired_seconds'),
  170. policy=profile.get('policy'),
  171. external_id=profile.get('external_id'),
  172. sts_region_id=profile.get('sts_region'),
  173. enable_vpc=profile.get('enable_vpc'),
  174. )
  175. elif mode == "CloudSSO":
  176. return CloudSSOCredentialsProvider(
  177. sign_in_url=profile.get('cloud_sso_sign_in_url'),
  178. account_id=profile.get('cloud_sso_account_id'),
  179. access_config=profile.get('cloud_sso_access_config'),
  180. access_token=profile.get('access_token'),
  181. access_token_expire=profile.get('cloud_sso_access_token_expire'),
  182. )
  183. elif mode == "OAuth":
  184. # 获取 OAuth 配置
  185. site_type = profile.get('oauth_site_type', 'CN')
  186. oauth_base_url_map = {
  187. 'CN': 'https://oauth.aliyun.com',
  188. 'INTL': 'https://oauth.alibabacloud.com'
  189. }
  190. sign_in_url = oauth_base_url_map.get(site_type.upper())
  191. if not sign_in_url:
  192. raise CredentialException('Invalid OAuth site type, support CN or INTL')
  193. oauth_client_map = {
  194. 'CN': '4038181954557748008',
  195. 'INTL': '4103531455503354461'
  196. }
  197. client_id = oauth_client_map.get(site_type.upper())
  198. if not client_id:
  199. raise CredentialException('Invalid OAuth site type, support CN or INTL')
  200. return OAuthCredentialsProvider(
  201. client_id=client_id,
  202. sign_in_url=sign_in_url,
  203. access_token=profile.get('oauth_access_token'),
  204. access_token_expire=profile.get('oauth_access_token_expire'),
  205. refresh_token=profile.get('oauth_refresh_token'),
  206. token_update_callback=self._get_oauth_token_update_callback(),
  207. token_update_callback_async=self._get_oauth_token_update_callback_async(),
  208. )
  209. elif mode == "External":
  210. return ExternalCredentialsProvider(
  211. process_command=profile.get('process_command'),
  212. credential_update_callback=self._get_external_credential_update_callback(),
  213. credential_update_callback_async=self._get_external_credential_update_callback_async(),
  214. )
  215. else:
  216. raise CredentialException(f"unsupported profile mode '{mode}' form cli credentials file.")
  217. raise CredentialException(f"unable to get profile with '{profile_name}' form cli credentials file.")
  218. def get_provider_name(self) -> str:
  219. return 'cli_profile'
  220. def _update_oauth_tokens(self, refresh_token: str, access_token: str, access_key: str, secret: str,
  221. security_token: str, access_token_expire: int, sts_expire: int) -> None:
  222. """更新 OAuth 令牌并写回配置文件"""
  223. def _find_source_oauth_profile(config: dict, profile_name: str) -> dict:
  224. profiles = config.get('profiles', [])
  225. profile = next((p for p in profiles if p.get('name') == profile_name), None)
  226. if not profile:
  227. raise CredentialException(f"unable to get profile with name '{profile_name}' from cli credentials file.")
  228. if profile.get('mode') == 'OAuth':
  229. return profile
  230. else:
  231. source_profile = profile.get('source_profile')
  232. if source_profile:
  233. return _find_source_oauth_profile(config, source_profile)
  234. raise CredentialException(f"unable to get OAuth profile with name '{profile_name}' from cli credentials file.")
  235. with self._file_lock:
  236. try:
  237. # 读取现有配置
  238. config = _load_config(self._profile_file)
  239. # 找到当前 profile 并更新 OAuth 令牌
  240. profile_name = self._profile_name or config.get('current')
  241. if not profile_name:
  242. raise CredentialException(f"unable to get profile to updated.")
  243. source_profile = _find_source_oauth_profile(config, profile_name)
  244. # 更新 OAuth 令牌
  245. source_profile['oauth_refresh_token'] = refresh_token
  246. source_profile['oauth_access_token'] = access_token
  247. source_profile['oauth_access_token_expire'] = access_token_expire
  248. # 更新 STS 凭据
  249. source_profile['access_key_id'] = access_key
  250. source_profile['access_key_secret'] = secret
  251. source_profile['sts_token'] = security_token
  252. source_profile['sts_expiration'] = sts_expire
  253. self._write_configuration_to_file_with_lock(self._profile_file, config)
  254. except Exception as e:
  255. raise CredentialException(f"failed to update OAuth tokens in config file: {e}")
  256. def _write_configuration_to_file(self, config_path: str, config: Dict) -> None:
  257. """将配置写入文件,使用原子写入确保数据完整性"""
  258. # 获取原文件权限(如果存在)
  259. file_mode = 0o644
  260. if os.path.exists(config_path):
  261. file_mode = os.stat(config_path).st_mode
  262. # 创建唯一临时文件
  263. import time
  264. temp_file = config_path + '.tmp-' + str(int(time.time() * 1000000)) # 微秒级时间戳
  265. backup_file = None
  266. try:
  267. # 写入临时文件
  268. self._write_config_file(temp_file, file_mode, config)
  269. # 原子性重命名,Windows下需要特殊处理
  270. if platform.system() == 'Windows' and self._allow_config_force_rewrite:
  271. # Windows下需要先删除目标文件,使用备份机制确保数据安全
  272. if os.path.exists(config_path):
  273. backup_file = config_path + '.backup'
  274. # 创建备份
  275. import shutil
  276. shutil.copy2(config_path, backup_file)
  277. # 删除原文件
  278. os.remove(config_path)
  279. os.rename(temp_file, config_path)
  280. # 成功后删除备份
  281. if backup_file and os.path.exists(backup_file):
  282. os.remove(backup_file)
  283. except Exception as e:
  284. # 恢复原文件(如果存在备份)
  285. if backup_file and os.path.exists(backup_file):
  286. try:
  287. if not os.path.exists(config_path):
  288. os.rename(backup_file, config_path)
  289. except Exception as restore_error:
  290. raise CredentialException(
  291. f"Failed to restore original file after write error: {restore_error}. Original error: {e}")
  292. raise e
  293. def _write_config_file(self, filename: str, file_mode: int, config: Dict) -> None:
  294. try:
  295. with open(filename, 'w', encoding='utf-8') as f:
  296. json.dump(config, f, indent=4, ensure_ascii=False)
  297. # 设置文件权限
  298. os.chmod(filename, file_mode)
  299. except Exception as e:
  300. raise CredentialException(f"Failed to write config file: {e}")
  301. def _write_configuration_to_file_with_lock(self, config_path: str, config: Dict) -> None:
  302. """使用操作系统级别的文件锁写入配置文件"""
  303. # 获取原文件权限(如果存在)
  304. file_mode = 0o644
  305. if os.path.exists(config_path):
  306. file_mode = os.stat(config_path).st_mode
  307. backup_file = None
  308. try:
  309. # 确保文件存在
  310. if not os.path.exists(config_path):
  311. # 创建空文件
  312. with open(config_path, 'w') as f:
  313. json.dump({}, f)
  314. # 在获取文件锁之前创建备份(Windows下需要)
  315. if platform.system() == 'Windows' and self._allow_config_force_rewrite and os.path.exists(config_path):
  316. backup_file = config_path + '.backup'
  317. import shutil
  318. shutil.copy2(config_path, backup_file)
  319. # 打开文件用于锁定
  320. with open(config_path, 'r+') as f:
  321. # 获取独占锁(阻塞其他进程)
  322. if HAS_MSVCRT:
  323. # Windows使用msvcrt
  324. msvcrt.locking(f.fileno(), msvcrt.LK_NBLCK, 1)
  325. elif HAS_FCNTL:
  326. # Unix/Linux使用fcntl
  327. fcntl.flock(f.fileno(), fcntl.LOCK_EX)
  328. # 如果都不支持,则跳过文件锁(仅进程内保护)
  329. try:
  330. if platform.system() == 'Windows' and self._allow_config_force_rewrite:
  331. # Windows下直接在锁定的文件中写入
  332. f.seek(0)
  333. f.truncate() # 清空文件内容
  334. json.dump(config, f, indent=4, ensure_ascii=False)
  335. f.flush()
  336. else:
  337. # 其他环境使用临时文件+rename(在文件锁内部进行原子操作)
  338. import time
  339. temp_file = config_path + '.tmp-' + str(int(time.time() * 1000000))
  340. self._write_config_file(temp_file, file_mode, config)
  341. # 在文件锁内部进行原子重命名
  342. os.rename(temp_file, config_path)
  343. finally:
  344. # 释放锁
  345. try:
  346. if HAS_MSVCRT:
  347. msvcrt.locking(f.fileno(), msvcrt.LK_UNLCK, 1)
  348. elif HAS_FCNTL:
  349. fcntl.flock(f.fileno(), fcntl.LOCK_UN)
  350. except (OSError, PermissionError):
  351. # 在Windows下,如果文件被重命名,文件句柄可能已经无效
  352. # 这种情况下锁会自动释放,所以忽略错误
  353. pass
  354. # 成功后删除备份
  355. if backup_file and os.path.exists(backup_file):
  356. os.remove(backup_file)
  357. except Exception as e:
  358. # 恢复原文件(如果存在备份)
  359. if backup_file and os.path.exists(backup_file):
  360. try:
  361. if not os.path.exists(config_path):
  362. os.rename(backup_file, config_path)
  363. except Exception as restore_error:
  364. raise CredentialException(
  365. f"Failed to restore original file after write error: {restore_error}. Original error: {e}")
  366. raise e
  367. def _get_oauth_token_update_callback(self) -> OAuthTokenUpdateCallback:
  368. """获取 OAuth 令牌更新回调函数"""
  369. return lambda refresh_token, access_token, access_key, secret, security_token, access_token_expire, sts_expire: self._update_oauth_tokens(
  370. refresh_token, access_token, access_key, secret, security_token, access_token_expire, sts_expire
  371. )
  372. def _update_external_credentials(self, access_key: str, secret: str,
  373. security_token: str, expiration: int) -> None:
  374. """更新 External 凭证并写回配置文件"""
  375. def _find_source_external_profile(config: dict, profile_name: str) -> dict:
  376. profiles = config.get('profiles', [])
  377. profile = next((p for p in profiles if p.get('name') == profile_name), None)
  378. if not profile:
  379. raise CredentialException(f"unable to get profile with name '{profile_name}' from cli credentials file.")
  380. if profile.get('mode') == 'External':
  381. return profile
  382. source_profile = profile.get('source_profile')
  383. if source_profile:
  384. return _find_source_external_profile(config, source_profile)
  385. raise CredentialException(f"unable to get External profile with name '{profile_name}' from cli credentials file.")
  386. with self._file_lock:
  387. try:
  388. config = _load_config(self._profile_file)
  389. profile_name = self._profile_name or config.get('current')
  390. if not profile_name:
  391. raise CredentialException(f"unable to get profile to updated.")
  392. source_profile = _find_source_external_profile(config, profile_name)
  393. source_profile['access_key_id'] = access_key
  394. source_profile['access_key_secret'] = secret
  395. source_profile['sts_token'] = security_token
  396. source_profile['sts_expiration'] = expiration
  397. self._write_configuration_to_file_with_lock(self._profile_file, config)
  398. except Exception as e:
  399. raise CredentialException(f"failed to update External credentials in config file: {e}")
  400. def _get_external_credential_update_callback(self) -> ExternalCredentialUpdateCallback:
  401. """获取 External 凭证更新回调函数"""
  402. return lambda access_key, secret, security_token, expiration: self._update_external_credentials(
  403. access_key, secret, security_token, expiration
  404. )
  405. async def _write_configuration_to_file_async(self, config_path: str, config: Dict) -> None:
  406. """异步将配置写入文件,使用原子写入确保数据完整性"""
  407. # 获取原文件权限(如果存在)
  408. file_mode = 0o644
  409. if os.path.exists(config_path):
  410. file_mode = os.stat(config_path).st_mode
  411. # 创建唯一临时文件
  412. import time
  413. temp_file = config_path + '.tmp-' + str(int(time.time() * 1000000)) # 微秒级时间戳
  414. backup_file = None
  415. try:
  416. # 异步写入临时文件
  417. await self._write_config_file_async(temp_file, file_mode, config)
  418. # 原子性重命名,Windows下需要特殊处理
  419. if platform.system() == 'Windows' and self._allow_config_force_rewrite:
  420. # Windows下需要先删除目标文件,使用备份机制确保数据安全
  421. if os.path.exists(config_path):
  422. backup_file = config_path + '.backup'
  423. # 创建备份
  424. import shutil
  425. shutil.copy2(config_path, backup_file)
  426. # 删除原文件
  427. os.remove(config_path)
  428. os.rename(temp_file, config_path)
  429. # 成功后删除备份
  430. if backup_file and os.path.exists(backup_file):
  431. os.remove(backup_file)
  432. except Exception as e:
  433. # 恢复原文件(如果存在备份)
  434. if backup_file and os.path.exists(backup_file):
  435. try:
  436. if not os.path.exists(config_path):
  437. os.rename(backup_file, config_path)
  438. except Exception as restore_error:
  439. raise CredentialException(
  440. f"Failed to restore original file after write error: {restore_error}. Original error: {e}")
  441. raise e
  442. async def _write_config_file_async(self, filename: str, file_mode: int, config: Dict) -> None:
  443. try:
  444. async with aiofiles.open(filename, 'w', encoding='utf-8') as f:
  445. await f.write(json.dumps(config, indent=4, ensure_ascii=False))
  446. # 设置文件权限
  447. os.chmod(filename, file_mode)
  448. except Exception as e:
  449. raise CredentialException(f"Failed to write config file: {e}")
  450. async def _write_configuration_to_file_with_lock_async(self, config_path: str, config: Dict) -> None:
  451. """异步使用操作系统级别的文件锁写入配置文件"""
  452. # 获取原文件权限(如果存在)
  453. file_mode = 0o644
  454. if os.path.exists(config_path):
  455. file_mode = os.stat(config_path).st_mode
  456. backup_file = None
  457. try:
  458. # 确保文件存在
  459. if not os.path.exists(config_path):
  460. # 创建空文件
  461. with open(config_path, 'w') as f:
  462. json.dump({}, f)
  463. # 在获取文件锁之前创建备份(Windows下需要)
  464. if platform.system() == 'Windows' and self._allow_config_force_rewrite and os.path.exists(config_path):
  465. backup_file = config_path + '.backup'
  466. import shutil
  467. shutil.copy2(config_path, backup_file)
  468. # 打开文件用于锁定
  469. with open(config_path, 'r+') as f:
  470. # 获取独占锁(阻塞其他进程)
  471. if HAS_MSVCRT:
  472. # Windows使用msvcrt
  473. msvcrt.locking(f.fileno(), msvcrt.LK_NBLCK, 1)
  474. elif HAS_FCNTL:
  475. # Unix/Linux使用fcntl
  476. fcntl.flock(f.fileno(), fcntl.LOCK_EX)
  477. # 如果都不支持,则跳过文件锁(仅进程内保护)
  478. try:
  479. if platform.system() == 'Windows' and self._allow_config_force_rewrite:
  480. # Windows下直接在锁定的文件中写入
  481. f.seek(0)
  482. f.truncate() # 清空文件内容
  483. json.dump(config, f, indent=4, ensure_ascii=False)
  484. f.flush()
  485. else:
  486. # 其他环境使用临时文件+rename(在文件锁内部进行原子操作)
  487. import time
  488. temp_file = config_path + '.tmp-' + str(int(time.time() * 1000000))
  489. await self._write_config_file_async(temp_file, file_mode, config)
  490. # 在文件锁内部进行原子重命名
  491. os.rename(temp_file, config_path)
  492. finally:
  493. # 释放锁
  494. try:
  495. if HAS_MSVCRT:
  496. msvcrt.locking(f.fileno(), msvcrt.LK_UNLCK, 1)
  497. elif HAS_FCNTL:
  498. fcntl.flock(f.fileno(), fcntl.LOCK_UN)
  499. except (OSError, PermissionError):
  500. # 在Windows下,如果文件被重命名,文件句柄可能已经无效
  501. # 这种情况下锁会自动释放,所以忽略错误
  502. pass
  503. # 成功后删除备份
  504. if backup_file and os.path.exists(backup_file):
  505. os.remove(backup_file)
  506. except Exception as e:
  507. # 恢复原文件(如果存在备份)
  508. if backup_file and os.path.exists(backup_file):
  509. try:
  510. if not os.path.exists(config_path):
  511. os.rename(backup_file, config_path)
  512. except Exception as restore_error:
  513. raise CredentialException(
  514. f"Failed to restore original file after write error: {restore_error}. Original error: {e}")
  515. raise e
  516. async def _update_oauth_tokens_async(self, refresh_token: str, access_token: str, access_key: str, secret: str,
  517. security_token: str, access_token_expire: int, sts_expire: int) -> None:
  518. """异步更新 OAuth 令牌并写回配置文件"""
  519. def _find_source_oauth_profile(config: dict, profile_name: str) -> dict:
  520. profiles = config.get('profiles', [])
  521. profile = next((p for p in profiles if p.get('name') == profile_name), None)
  522. if not profile:
  523. raise CredentialException(f"unable to get profile with name '{profile_name}' from cli credentials file.")
  524. if profile.get('mode') == 'OAuth':
  525. return profile
  526. else:
  527. source_profile = profile.get('source_profile')
  528. if source_profile:
  529. return _find_source_oauth_profile(config, source_profile)
  530. raise CredentialException(f"unable to get OAuth profile with name '{profile_name}' from cli credentials file.")
  531. with self._file_lock:
  532. try:
  533. # 读取现有配置
  534. config = await _load_config_async(self._profile_file)
  535. # 找到当前 profile 并更新 OAuth 令牌
  536. profile_name = self._profile_name or config.get('current')
  537. if not profile_name:
  538. raise CredentialException(f"unable to get profile to updated.")
  539. source_profile = _find_source_oauth_profile(config, profile_name)
  540. # 更新 OAuth 令牌
  541. source_profile['oauth_refresh_token'] = refresh_token
  542. source_profile['oauth_access_token'] = access_token
  543. source_profile['oauth_access_token_expire'] = access_token_expire
  544. # 更新 STS 凭据
  545. source_profile['access_key_id'] = access_key
  546. source_profile['access_key_secret'] = secret
  547. source_profile['sts_token'] = security_token
  548. source_profile['sts_expiration'] = sts_expire
  549. await self._write_configuration_to_file_with_lock_async(self._profile_file, config)
  550. except Exception as e:
  551. raise CredentialException(f"failed to update OAuth tokens in config file: {e}")
  552. def _get_oauth_token_update_callback_async(self) -> OAuthTokenUpdateCallbackAsync:
  553. """获取异步 OAuth 令牌更新回调函数"""
  554. return lambda refresh_token, access_token, access_key, secret, security_token, access_token_expire, sts_expire: self._update_oauth_tokens_async(
  555. refresh_token, access_token, access_key, secret, security_token, access_token_expire, sts_expire
  556. )
  557. async def _update_external_credentials_async(self, access_key: str, secret: str,
  558. security_token: str, expiration: int) -> None:
  559. """异步更新 External 凭证并写回配置文件"""
  560. def _find_source_external_profile(config: dict, profile_name: str) -> dict:
  561. profiles = config.get('profiles', [])
  562. profile = next((p for p in profiles if p.get('name') == profile_name), None)
  563. if not profile:
  564. raise CredentialException(f"unable to get profile with name '{profile_name}' from cli credentials file.")
  565. if profile.get('mode') == 'External':
  566. return profile
  567. source_profile = profile.get('source_profile')
  568. if source_profile:
  569. return _find_source_external_profile(config, source_profile)
  570. raise CredentialException(f"unable to get External profile with name '{profile_name}' from cli credentials file.")
  571. with self._file_lock:
  572. try:
  573. config = await _load_config_async(self._profile_file)
  574. profile_name = self._profile_name or config.get('current')
  575. if not profile_name:
  576. raise CredentialException(f"unable to get profile to updated.")
  577. source_profile = _find_source_external_profile(config, profile_name)
  578. source_profile['access_key_id'] = access_key
  579. source_profile['access_key_secret'] = secret
  580. source_profile['sts_token'] = security_token
  581. source_profile['sts_expiration'] = expiration
  582. await self._write_configuration_to_file_with_lock_async(self._profile_file, config)
  583. except Exception as e:
  584. raise CredentialException(f"failed to update External credentials in config file: {e}")
  585. def _get_external_credential_update_callback_async(self) -> ExternalCredentialUpdateCallbackAsync:
  586. """获取异步 External 凭证更新回调函数"""
  587. return lambda access_key, secret, security_token, expiration: self._update_external_credentials_async(
  588. access_key, secret, security_token, expiration
  589. )