Merge pull request #205 from zeyu-zh/master

支持credentials provider

README.md

@@ -56,6 +56,7 @@ client = NacosClient(server_addresses, namespace=your_ns, ak=your_ak, sk=your_sk
 * *namespace* - Namespace. | default: `None`
 * *ak* - The accessKey to authenticate. | default: null
 * *sk* - The secretKey to authentication. | default: null
+* *credentials_provider* - The custom access key manager | default: null
 * *log_level* - Log level. | default: null
 * *log_rotation_backup_count* - The number of log files to keep. | default: `7`
 
@@ -314,6 +315,7 @@ client_config = (ClientConfigBuilder()
 * *server_address* - **required**  - Nacos server address
 * *access_key* - The aliyun accessKey to authenticate.
 * *secret_key* - The aliyun secretKey to authenticate.
+* *credentials_provider* - The custom access key manager.
 * *username* - The username to authenticate.
 * *password* - The password to authenticate.
 * *log_level* - Log level | default: `logging.INFO`

nacos/auth.py

@@ -0,0 +1,26 @@
+class Credentials(object):
+    def __init__(self, access_key_id, access_key_secret, security_token=None):
+        self.access_key_id = access_key_id
+        self.access_key_secret = access_key_secret
+        self.security_token = security_token
+
+    def get_access_key_id(self):
+        return self.access_key_id
+
+    def get_access_key_secret(self):
+        return self.access_key_secret
+
+    def get_security_token(self):
+        return self.security_token
+
+class CredentialsProvider(object):
+    def get_credentials(self):
+        return
+
+
+class StaticCredentialsProvider(CredentialsProvider):
+    def __init__(self, access_key_id="", access_key_secret="", security_token=""):
+        self.credentials = Credentials(access_key_id, access_key_secret, security_token)
+
+    def get_credentials(self):
+        return self.credentials

nacos/client.py

@@ -14,6 +14,7 @@
 from typing import Dict
 
 from .task import HeartbeatInfo, HeartbeatTask
+from .auth import StaticCredentialsProvider
 
 try:
     import ssl
@@ -276,7 +277,8 @@ def initLog(self, logDir, log_level, log_rotation_backup_count):
             logger.propagate = False
 
     def __init__(self, server_addresses=None, endpoint=None, namespace=None, ak=None,
-                 sk=None, username=None, password=None, logDir=None, log_level=None, log_rotation_backup_count=None):
+                 sk=None, username=None, password=None, logDir=None, log_level=None,
+                 log_rotation_backup_count=None, credentials_provider=None):
         self.server_list = list()
         self.initLog(logDir, log_level, log_rotation_backup_count)
         try:
@@ -308,8 +310,7 @@ def __init__(self, server_addresses=None, endpoint=None, namespace=None, ak=None
 
         self.endpoint = endpoint
         self.namespace = namespace or DEFAULT_NAMESPACE or ""
-        self.ak = ak
-        self.sk = sk
+        self.credentials_provider = credentials_provider if credentials_provider else StaticCredentialsProvider(ak, sk)
         self.username = username
         self.password = password
 
@@ -330,7 +331,8 @@ def __init__(self, server_addresses=None, endpoint=None, namespace=None, ak=None
         self.process_mgr = None
 
         self.default_timeout = DEFAULTS["TIMEOUT"]
-        self.auth_enabled = self.ak and self.sk
+        credentials = self.credentials_provider.get_credentials()
+        self.auth_enabled = credentials.get_access_key_id() and credentials.get_access_key_secret()
         self.cai_enabled = True
         self.pulling_timeout = DEFAULTS["PULLING_TIMEOUT"]
         self.pulling_config_size = DEFAULTS["PULLING_CONFIG_SIZE"]
@@ -884,15 +886,16 @@ def _inject_auth_info(self, headers, params, data, module="config"):
         if not params and not data:
             return
         ts = str(int(time.time() * 1000))
-        ak, sk = self.ak, self.sk
+        # now we have a fixed credentials (access key or sts token)
+        credentials = self.credentials_provider.get_credentials()
 
         sign_str = ""
 
         params_to_sign = params or data or {}
         # config signature
         if "config" == module:
             headers.update({
-                "Spas-AccessKey": ak,
+                "Spas-AccessKey": credentials.get_access_key_id(),
                 "timeStamp": ts,
             })
 
@@ -905,7 +908,9 @@ def _inject_auth_info(self, headers, params, data, module="config"):
                 sign_str = sign_str + group + "+"
             if sign_str:
                 sign_str += ts
-                headers["Spas-Signature"] = self.__do_sign(sign_str, sk)
+                headers["Spas-Signature"] = self.__do_sign(sign_str, credentials.get_access_key_secret())
+            if credentials.get_security_token():
+                headers["Spas-SecurityToken"] = credentials.get_security_token()
 
         # naming signature
         else:
@@ -922,10 +927,12 @@ def _inject_auth_info(self, headers, params, data, module="config"):
                 sign_str = ts
 
             params.update({
-                "ak": ak,
+                "ak": credentials.get_access_key_id(),
                 "data": sign_str,
-                "signature": self.__do_sign(sign_str, sk),
+                "signature": self.__do_sign(sign_str, credentials.get_access_key_secret()),
             })
+            if credentials.get_security_token():
+                params.update({"Spas-SecurityToken": credentials.get_security_token()})
 
     def __do_sign(self, sign_str, sk):
         return base64.encodebytes(

test/client_test.py

@@ -8,6 +8,7 @@
 from nacos import files
 from nacos.listener import SubscribeListener, SimpleListenerManager
 from nacos.timer import NacosTimer, NacosTimerManager
+from nacos.auth import CredentialsProvider, Credentials
 import time
 import shutil
 
@@ -25,6 +26,13 @@
 # Set the following option if http requests need through by proxy
 # client.set_options(proxies={"http":"192.168.56.1:809"})
 
+class CustomCredentialsProvider(CredentialsProvider):
+    def __init__(self, ak="", sk="", token=""):
+        self.credential = Credentials(ak, sk, token)
+
+    def get_credentials(self):
+        return self.credential
+
 class TestClient(unittest.TestCase):
     def test_get_server(self):
         self.assertEqual(client.get_server(), (SERVER_1, 8848))
@@ -315,7 +323,7 @@ def test_auth_off(self):
         self.assertEqual(0, len(headers))
         self.assertEqual(0, len(params))
 
-    def test_inject_auth_info_of_config(self):
+    def test_inject_auth_info_of_config_with_fixed_ak(self):
         headers = {}
         params = {"tenant": "abc", "group": "bbb"}
         client_auth = nacos.NacosClient(SERVER_ADDRESSES, ak="1", sk="1")
@@ -325,7 +333,7 @@ def test_inject_auth_info_of_config(self):
         self.assertTrue("timeStamp" in headers)
         self.assertTrue("Spas-Signature" in headers)
 
-    def test_inject_auth_info_of_naming(self):
+    def test_inject_auth_info_of_naming_with_fixed_ak(self):
         headers = {}
         params = {"serviceName": "abc", "groupName": "bbb"}
         client_auth = nacos.NacosClient(SERVER_ADDRESSES, ak="1", sk="1")
@@ -335,6 +343,48 @@ def test_inject_auth_info_of_naming(self):
         self.assertTrue("data" in params)
         self.assertTrue("signature" in params)
 
+    def test_inject_auth_info_of_config_with_provider(self):
+        headers = {}
+        params = {"tenant": "abc", "group": "bbb"}
+
+        # access_key
+        client_auth = nacos.NacosClient(SERVER_ADDRESSES, credentials_provider=CustomCredentialsProvider("1", "1"))
+        self.assertFalse(client_auth.auth_enabled is None)
+        client_auth._inject_auth_info(headers, params, data=None, module="config")
+        self.assertEqual("1", headers.get("Spas-AccessKey"))
+        self.assertTrue("timeStamp" in headers)
+        self.assertTrue("Spas-Signature" in headers)
+        self.assertTrue("Spas-SecurityToken" not in headers)
+
+        # security_token
+        client_auth = nacos.NacosClient(SERVER_ADDRESSES, credentials_provider=CustomCredentialsProvider("1", "1", "1"))
+        self.assertFalse(client_auth.auth_enabled is None)
+        client_auth._inject_auth_info(headers, params, data=None, module="config")
+        self.assertEqual("1", headers.get("Spas-AccessKey"))
+        self.assertEqual("1", headers.get("Spas-SecurityToken"))
+        self.assertTrue("timeStamp" in headers)
+        self.assertTrue("Spas-Signature" in headers)
+
+    def test_inject_auth_info_of_naming_with_provider(self):
+        headers = {}
+        params = {"serviceName": "abc", "groupName": "bbb"}
+        # access_key
+        client_auth = nacos.NacosClient(SERVER_ADDRESSES, credentials_provider=CustomCredentialsProvider("1", "1"))
+        self.assertFalse(client_auth.auth_enabled is None)
+        client_auth._inject_auth_info(headers, params, data=None, module="naming")
+        self.assertEqual("1", params.get("ak"))
+        self.assertTrue("data" in params)
+        self.assertTrue("signature" in params)
+        self.assertTrue("Spas-SecurityToken" not in headers)
+
+        # security_token
+        client_auth = nacos.NacosClient(SERVER_ADDRESSES, credentials_provider=CustomCredentialsProvider("1", "1", "1"))
+        self.assertFalse(client_auth.auth_enabled is None)
+        client_auth._inject_auth_info(headers, params, data=None, module="naming")
+        self.assertEqual("1", params.get("ak"))
+        self.assertEqual("1", params.get("Spas-SecurityToken"))
+        self.assertTrue("data" in params)
+        self.assertTrue("signature" in params)
 
 if __name__ == '__main__':
     unittest.main()

test/client_v2_test.py

@@ -8,6 +8,7 @@
 from v2.nacos.naming.model.naming_param import RegisterInstanceParam, DeregisterInstanceParam, \
     BatchRegisterInstanceParam, GetServiceParam, ListServiceParam, SubscribeServiceParam, ListInstanceParam
 from v2.nacos.naming.nacos_naming_service import NacosNamingService
+from v2.nacos.common.auth import CredentialsProvider, Credentials
 
 client_config = (ClientConfigBuilder()
                  .access_key(os.getenv('NACOS_ACCESS_KEY'))
@@ -17,10 +18,16 @@
                  .grpc_config(GRPCConfig(grpc_timeout=5000))
                  .build())
 
+class CustomCredentialsProvider(CredentialsProvider):
+    def __init__(self, ak="", sk="", token=""):
+        self.credential = Credentials(ak, sk, token)
+
+    def get_credentials(self):
+        return self.credential
 
 class TestClientV2(unittest.IsolatedAsyncioTestCase):
 
-    async def test_register_with_endpoint(self):
+    async def test_register_with_endpoint_and_fixed_ak(self):
         config = (ClientConfigBuilder()
                          .access_key(os.getenv('NACOS_ACCESS_KEY'))
                          .secret_key(os.getenv('NACOS_SECRET_KEY'))
@@ -32,6 +39,17 @@ async def test_register_with_endpoint(self):
         client = await NacosNamingService.create_naming_service(config)
         assert await client.server_health()
 
+    async def test_register_with_endpoint_and_provider(self):
+        config = (ClientConfigBuilder()
+                         .credentials_provider(CustomCredentialsProvider(os.getenv('NACOS_ACCESS_KEY'), os.getenv('NACOS_SECRET_KEY')))
+                         .endpoint(os.getenv('NACOS_SERVER_ENDPOINT', 'localhost:8848'))
+                         .log_level('INFO')
+                         .grpc_config(GRPCConfig(grpc_timeout=5000))
+                         .build())
+
+        client = await NacosNamingService.create_naming_service(config)
+        assert await client.server_health()
+
     async def test_register(self):
         client = await NacosNamingService.create_naming_service(client_config)
         assert await client.server_health()

test/config_client_v2_test.py

@@ -6,6 +6,7 @@
 from v2.nacos.common.client_config_builder import ClientConfigBuilder
 from v2.nacos.config.model.config_param import ConfigParam
 from v2.nacos.config.nacos_config_service import NacosConfigService
+from v2.nacos.common.auth import CredentialsProvider, Credentials
 
 client_config = (ClientConfigBuilder()
                  .access_key(os.getenv('NACOS_ACCESS_KEY'))
@@ -15,6 +16,12 @@
                  .grpc_config(GRPCConfig(grpc_timeout=5000))
                  .build())
 
+class CustomCredentialsProvider(CredentialsProvider):
+    def __init__(self, ak="", sk="", token=""):
+        self.credential = Credentials(ak, sk, token)
+
+    def get_credentials(self):
+        return self.credential
 
 class TestClientV2(unittest.IsolatedAsyncioTestCase):
     async def test_publish_config(self):
@@ -235,15 +242,7 @@ async def config_listener(tenant, data_id, group, content):
 
         await client.shutdown()
 
-    async def test_gray_config(self):
-        client_cfg = (ClientConfigBuilder()
-                      .access_key(os.getenv('NACOS_ACCESS_KEY'))
-                      .secret_key(os.getenv('NACOS_SECRET_KEY'))
-                      .server_address(os.getenv('NACOS_SERVER_ADDR', 'localhost:8848'))
-                      .log_level('INFO')
-                      .app_conn_labels({"k1": "v1", "k2": "v2", "nacos_config_gray_label": "gray"})
-                      .grpc_config(GRPCConfig(grpc_timeout=5000))
-                      .build())
+    async def _gray_config(self, client_cfg):
         client = await NacosConfigService.create_config_service(client_cfg)
 
         dataID = "com.alibaba.nacos.test.config.gray"
@@ -255,3 +254,24 @@ async def config_listener(tenant, data_id, group, content):
         await client.add_listener(dataID, groupName, config_listener)
 
         await asyncio.sleep(1000)
+
+    async def test_gray_config_with_fixed_ak(self):
+        client_cfg = (ClientConfigBuilder()
+                      .access_key(os.getenv('NACOS_ACCESS_KEY'))
+                      .secret_key(os.getenv('NACOS_SECRET_KEY'))
+                      .server_address(os.getenv('NACOS_SERVER_ADDR', 'localhost:8848'))
+                      .log_level('INFO')
+                      .app_conn_labels({"k1": "v1", "k2": "v2", "nacos_config_gray_label": "gray"})
+                      .grpc_config(GRPCConfig(grpc_timeout=5000))
+                      .build())
+        await self._gray_config(client_cfg)
+
+    async def test_gray_config_with_provider(self):
+        client_cfg = (ClientConfigBuilder()
+                      .credentials_provider(CustomCredentialsProvider(os.getenv('NACOS_ACCESS_KEY'), os.getenv('NACOS_SECRET_KEY')))
+                      .server_address(os.getenv('NACOS_SERVER_ADDR', 'localhost:8848'))
+                      .log_level('INFO')
+                      .app_conn_labels({"k1": "v1", "k2": "v2", "nacos_config_gray_label": "gray"})
+                      .grpc_config(GRPCConfig(grpc_timeout=5000))
+                      .build())
+        await self._gray_config(client_cfg)

v2/nacos/common/auth.py

@@ -0,0 +1,32 @@
+class Credentials(object):
+    def __init__(self, access_key_id, access_key_secret, security_token=None):
+        self.access_key_id = access_key_id
+        self.access_key_secret = access_key_secret
+        self.security_token = security_token
+
+    def get_access_key_id(self):
+        return self.access_key_id
+
+    def get_access_key_secret(self):
+        return self.access_key_secret
+
+    def get_security_token(self):
+        return self.security_token
+
+class CredentialsProvider(object):
+    def get_credentials(self):
+        return
+
+
+class StaticCredentialsProvider(CredentialsProvider):
+    def __init__(self, access_key_id="", access_key_secret="", security_token=""):
+        self.credentials = Credentials(access_key_id, access_key_secret, security_token)
+
+    def get_credentials(self):
+        return self.credentials
+
+    def set_access_key_id(self, access_key_id):
+        self.credentials.access_key_id = access_key_id
+
+    def set_access_key_secret(self, access_key_secret):
+        self.credentials.access_key_secret = access_key_secret

v2/nacos/common/client_config.py

@@ -2,6 +2,7 @@
 
 from v2.nacos.common.constants import Constants
 from v2.nacos.common.nacos_exception import NacosException, INVALID_PARAM
+from v2.nacos.common.auth import StaticCredentialsProvider
 
 
 class KMSConfig:
@@ -44,7 +45,7 @@ def __init__(self, max_receive_message_length=Constants.GRPC_MAX_RECEIVE_MESSAGE
 class ClientConfig:
     def __init__(self, server_addresses=None, endpoint=None, namespace_id='', context_path='', access_key=None,
                  secret_key=None, username=None, password=None, app_name='', app_key='', log_dir='', log_level=None,
-                 log_rotation_backup_count=None, app_conn_labels=None):
+                 log_rotation_backup_count=None, app_conn_labels=None, credentials_provider=None):
         self.server_list = []
         try:
             if server_addresses is not None and server_addresses.strip() != "":
@@ -57,9 +58,8 @@ def __init__(self, server_addresses=None, endpoint=None, namespace_id='', contex
         self.endpoint_context_path = Constants.WEB_CONTEXT
         self.endpoint_query_header = None
         self.namespace_id = namespace_id
-        self.access_key = access_key
+        self.credentials_provider = credentials_provider if credentials_provider else StaticCredentialsProvider(access_key, secret_key)
         self.context_path = context_path
-        self.secret_key = secret_key
         self.username = username  # the username for nacos auth
         self.password = password  # the password for nacos auth
         self.app_name = app_name