From c52daf8b66cd3f2d330f189b2719bff5256d5f9c Mon Sep 17 00:00:00 2001
From: Sonny <theli@outlook.com>
Date: Wed, 30 Nov 2022 17:04:52 -0500
Subject: [PATCH] tencentcloud - client - support for assume role (#8043)

---
 .../c7n_tencentcloud/client.py                | 26 +++++++++++-
 .../c7n_tencentcloud/tests/test_tc_client.py  | 40 ++++++++++++++++++-
 2 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/tools/c7n_tencentcloud/c7n_tencentcloud/client.py b/tools/c7n_tencentcloud/c7n_tencentcloud/client.py
index c701bd8c0ba..d3d1c36678f 100644
--- a/tools/c7n_tencentcloud/c7n_tencentcloud/client.py
+++ b/tools/c7n_tencentcloud/c7n_tencentcloud/client.py
@@ -1,6 +1,8 @@
 # Copyright The Cloud Custodian Authors.
 # SPDX-License-Identifier: Apache-2.0
 
+import os
+
 import jmespath
 import socket
 from retrying import retry
@@ -8,6 +10,7 @@
 from c7n.exceptions import PolicyExecutionError
 from requests.exceptions import ConnectionError
 from tencentcloud.common import credential
+from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
 from tencentcloud.common.profile.client_profile import ClientProfile
 from tencentcloud.common.profile.http_profile import HttpProfile
 from tencentcloud.common.common_client import CommonClient
@@ -124,7 +127,28 @@ def __init__(self) -> None:
         # just using default get_credentials() method
         # steps: Environment Variable -> profile file -> CVM role
         # for reference: https://github.com/TencentCloud/tencentcloud-sdk-python
-        self._cred = credential.DefaultCredentialProvider().get_credentials()
+
+        cred_provider = credential.DefaultCredentialProvider()
+
+        # the DefaultCredentialProvider does not handle sts assumed role sessions
+        # so we need to check for the token first
+        if 'TENCENTCLOUD_TOKEN' in os.environ:
+            if (
+                'TENCENTCLOUD_SECRET_ID' not in os.environ or
+                'TENCENTCLOUD_SECRET_KEY' not in os.environ
+            ):
+                raise TencentCloudSDKException(
+                    'TENCENTCLOUD_TOKEN provided, but one of TENCENTCLOUD_SECRET_ID'
+                    'or TENCENTCLOUD_SECRET_KEY missing'
+                )
+            cred = credential.Credential(
+                secret_id=os.environ['TENCENTCLOUD_SECRET_ID'],
+                secret_key=os.environ['TENCENTCLOUD_SECRET_KEY'],
+                token=os.environ['TENCENTCLOUD_TOKEN']
+            )
+            cred_provider.cred = cred
+
+        self._cred = cred_provider.get_credentials()
 
     def client(self,
                endpoint: str,
diff --git a/tools/c7n_tencentcloud/tests/test_tc_client.py b/tools/c7n_tencentcloud/tests/test_tc_client.py
index 0073fece4f9..627ebc3e228 100644
--- a/tools/c7n_tencentcloud/tests/test_tc_client.py
+++ b/tools/c7n_tencentcloud/tests/test_tc_client.py
@@ -1,12 +1,19 @@
 # Copyright The Cloud Custodian Authors.
 # SPDX-License-Identifier: Apache-2.0
 
+import os
+import socket
+
 from datetime import datetime
+from unittest.mock import patch
+
 import jmespath
 import pytest
-import socket
-from retrying import RetryError
+
 from c7n_tencentcloud.utils import PageMethod
+from c7n_tencentcloud.client import Session
+
+from retrying import RetryError
 from tencentcloud.common.abstract_client import AbstractClient
 from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
 
@@ -133,3 +140,32 @@ def test_paging_token(self, client_tag):
         }
         res = client_tag.execute_paged_query("GetTagValues", params, jsonpath, paging_def)
         assert len(res) == 233
+
+    @patch.dict(
+        os.environ,
+        {
+            "TENCENTCLOUD_TOKEN": "foo",
+            "TENCENTCLOUD_SECRET_KEY": "bar",
+            "TENCENTCLOUD_SECRET_ID": "baz",
+        }, clear=True
+    )
+    def test_tc_client_token(self):
+        session = Session()
+        assert session._cred.token == 'foo'
+        assert session._cred.secret_key == 'bar'
+        assert session._cred.secret_id == 'baz'
+
+    @patch.dict(
+        os.environ,
+        {
+            "TENCENTCLOUD_TOKEN": "foo",
+            "TENCENTCLOUD_SECRET_ID": "baz",
+        }, clear=True
+    )
+    def test_tc_client_token_missing_key(self):
+        found = False
+        try:
+            Session()
+        except TencentCloudSDKException:
+            found = True
+        assert found