ecs ram role credentials provider

liyanzhang505 · huiguangjun · commit ad6805d0e2c7 · 2021-02-09T09:10:24.000+08:00
diff --git a/oss2/__init__.py b/oss2/__init__.py
@@ -3,9 +3,9 @@
 from . import models, exceptions, defaults
 
 from .api import Service, Bucket
-from .auth import Auth, AuthV2, AnonymousAuth, StsAuth, AUTH_VERSION_1, AUTH_VERSION_2, make_auth
+from .auth import Auth, AuthV2, AnonymousAuth, StsAuth, AUTH_VERSION_1, AUTH_VERSION_2, make_auth, ProviderAuth, ProviderAuthV2
 from .http import Session, CaseInsensitiveDict
-
+from .credentials import EcsRamRoleCredentialsProvider, EcsRamRoleCredential, CredentialsProvider, StaticCredentialsProvider
 
 from .iterators import (BucketIterator, ObjectIterator, ObjectIteratorV2,
                         MultipartUploadIterator, ObjectUploadIterator,
diff --git a/oss2/auth.py b/oss2/auth.py
@@ -8,6 +8,7 @@
 from .compat import urlquote, to_bytes, is_py2
 from .headers import *
 import logging
+from .credentials import StaticCredentialsProvider
 
 AUTH_VERSION_1 = 'v1'
 AUTH_VERSION_2 = 'v2'
@@ -26,11 +27,14 @@ def make_auth(access_key_id, access_key_secret, auth_version=AUTH_VERSION_1):
 
 class AuthBase(object):
     """用于保存用户AccessKeyId、AccessKeySecret，以及计算签名的对象。"""
-    def __init__(self, access_key_id, access_key_secret):
-        self.id = access_key_id.strip()
-        self.secret = access_key_secret.strip()
+    def __init__(self, credentials_provider):
+        self.credentials_provider = credentials_provider
 
     def _sign_rtmp_url(self, url, bucket_name, channel_name, expires, params):
+        credentials = self.credentials_provider.get_credentials()
+        if credentials.get_security_token():
+            params['security-token'] = credentials.get_security_token()
+
         expiration_time = int(time.time()) + expires
 
         canonicalized_resource = "/%s/%s" % (bucket_name, channel_name)
@@ -52,18 +56,21 @@ def _sign_rtmp_url(self, url, bucket_name, channel_name, expires, params):
         logger.debug('Sign Rtmp url: string to be signed = {0}'.format(string_to_sign))
 
 
-        h = hmac.new(to_bytes(self.secret), to_bytes(string_to_sign), hashlib.sha1)
+        h = hmac.new(to_bytes(credentials.get_access_key_secret()), to_bytes(string_to_sign), hashlib.sha1)
         signature = utils.b64encode_as_string(h.digest())
 
-        p['OSSAccessKeyId'] = self.id
+        p['OSSAccessKeyId'] = credentials.get_access_key_id()
         p['Expires'] = str(expiration_time)
         p['Signature'] = signature
 
         return url + '?' + '&'.join(_param_to_quoted_query(k, v) for k, v in p.items())
 
 
-class Auth(AuthBase):
-    """签名版本1"""
+
+class ProviderAuth(AuthBase):
+    """签名版本1
+    默认构造函数同父类AuthBase，需要传递credentials_provider
+    """
     _subresource_key_set = frozenset(
         ['response-content-type', 'response-content-language',
          'response-cache-control', 'logging', 'response-content-encoding',
@@ -80,32 +87,40 @@ class Auth(AuthBase):
     )
 
     def _sign_request(self, req, bucket_name, key):
+        credentials = self.credentials_provider.get_credentials()
+        if credentials.get_security_token():
+            req.headers[OSS_SECURITY_TOKEN] = credentials.get_security_token()
+
         req.headers['date'] = utils.http_date()
 
-        signature = self.__make_signature(req, bucket_name, key)
-        req.headers['authorization'] = "OSS {0}:{1}".format(self.id, signature)
+        signature = self.__make_signature(req, bucket_name, key, credentials)
+        req.headers['authorization'] = "OSS {0}:{1}".format(credentials.get_access_key_id(), signature)
 
     def _sign_url(self, req, bucket_name, key, expires):
+        credentials = self.credentials_provider.get_credentials()
+        if credentials.get_security_token():
+            req.params['security-token'] = credentials.get_security_token()
+
         expiration_time = int(time.time()) + expires
 
         req.headers['date'] = str(expiration_time)
-        signature = self.__make_signature(req, bucket_name, key)
+        signature = self.__make_signature(req, bucket_name, key, credentials)
 
-        req.params['OSSAccessKeyId'] = self.id
+        req.params['OSSAccessKeyId'] = credentials.get_access_key_id()
         req.params['Expires'] = str(expiration_time)
         req.params['Signature'] = signature
 
         return req.url + '?' + '&'.join(_param_to_quoted_query(k, v) for k, v in req.params.items())
 
-    def __make_signature(self, req, bucket_name, key):
+    def __make_signature(self, req, bucket_name, key, credentials):
         if is_py2:
             string_to_sign = self.__get_string_to_sign(req, bucket_name, key)
         else:
             string_to_sign = self.__get_bytes_to_sign(req, bucket_name, key)
 
         logger.debug('Make signature: string to be signed = {0}'.format(string_to_sign))
 
-        h = hmac.new(to_bytes(self.secret), to_bytes(string_to_sign), hashlib.sha1)
+        h = hmac.new(to_bytes(credentials.get_access_key_secret()), to_bytes(string_to_sign), hashlib.sha1)
         return utils.b64encode_as_string(h.digest())
 
     def __get_string_to_sign(self, req, bucket_name, key):
@@ -192,6 +207,14 @@ def __get_headers_bytes(self, req):
         else:
             return b''
 
+class Auth(ProviderAuth):
+    """签名版本1
+    """
+    def __init__(self, access_key_id, access_key_secret):
+        credentials_provider = StaticCredentialsProvider(access_key_id.strip(), access_key_secret.strip())
+        super(Auth, self).__init__(credentials_provider)
+
+
 class AnonymousAuth(object):
     """用于匿名访问。
 
@@ -222,19 +245,16 @@ class StsAuth(object):
     """
     def __init__(self, access_key_id, access_key_secret, security_token, auth_version=AUTH_VERSION_1):
         logger.debug("Init StsAuth: access_key_id: {0}, access_key_secret: ******, security_token: ******".format(access_key_id))
-        self.__auth = make_auth(access_key_id, access_key_secret, auth_version)
-        self.__security_token = security_token
+        credentials_provider = StaticCredentialsProvider(access_key_id, access_key_secret, security_token)
+        self.__auth = ProviderAuthV2(credentials_provider) if auth_version == AUTH_VERSION_2 else ProviderAuth(credentials_provider)
 
     def _sign_request(self, req, bucket_name, key):
-        req.headers[OSS_SECURITY_TOKEN] = self.__security_token
         self.__auth._sign_request(req, bucket_name, key)
 
     def _sign_url(self, req, bucket_name, key, expires):
-        req.params['security-token'] = self.__security_token
         return self.__auth._sign_url(req, bucket_name, key, expires)
 
     def _sign_rtmp_url(self, url, bucket_name, channel_name, expires, params):
-        params['security-token'] = self.__security_token
         return self.__auth._sign_rtmp_url(url, bucket_name, channel_name, expires, params)
 
 
@@ -268,8 +288,9 @@ def v2_uri_encode(raw_text):
                                    'if-modified-since'])
 
 
-class AuthV2(AuthBase):
-    """签名版本2，与版本1的区别在：
+class ProviderAuthV2(AuthBase):
+    """签名版本2，默认构造函数同父类AuthBase，需要传递credentials_provider
+    与版本1的区别在：
     1. 使用SHA256算法，具有更高的安全性
     2. 参数计算包含所有的HTTP查询参数
     """
@@ -283,20 +304,24 @@ def _sign_request(self, req, bucket_name, key, in_additional_headers=None):
         :param key: OSS文件名
         :param in_additional_headers: 加入签名计算的额外header列表
         """
+        credentials = self.credentials_provider.get_credentials()
+        if credentials.get_security_token():
+            req.headers[OSS_SECURITY_TOKEN] = credentials.get_security_token()
+
         if in_additional_headers is None:
             in_additional_headers = _DEFAULT_ADDITIONAL_HEADERS
 
         additional_headers = self.__get_additional_headers(req, in_additional_headers)
 
         req.headers['date'] = utils.http_date()
 
-        signature = self.__make_signature(req, bucket_name, key, additional_headers)
+        signature = self.__make_signature(req, bucket_name, key, additional_headers, credentials)
 
         if additional_headers:
             req.headers['authorization'] = "OSS2 AccessKeyId:{0},AdditionalHeaders:{1},Signature:{2}"\
-                .format(self.id, ';'.join(additional_headers), signature)
+                .format(credentials.get_access_key_id(), ';'.join(additional_headers), signature)
         else:
-            req.headers['authorization'] = "OSS2 AccessKeyId:{0},Signature:{1}".format(self.id, signature)
+            req.headers['authorization'] = "OSS2 AccessKeyId:{0},Signature:{1}".format(credentials.get_access_key_id(), signature)
 
     def _sign_url(self, req, bucket_name, key, expires, in_additional_headers=None):
         """返回一个签过名的URL
@@ -311,6 +336,9 @@ def _sign_url(self, req, bucket_name, key, expires, in_additional_headers=None):
 
         :return: a signed URL
         """
+        credentials = self.credentials_provider.get_credentials()
+        if credentials.get_security_token():
+            req.params['security-token'] = credentials.get_security_token()
 
         if in_additional_headers is None:
             in_additional_headers = set()
@@ -323,23 +351,23 @@ def _sign_url(self, req, bucket_name, key, expires, in_additional_headers=None):
 
         req.params['x-oss-signature-version'] = 'OSS2'
         req.params['x-oss-expires'] = str(expiration_time)
-        req.params['x-oss-access-key-id'] = self.id
+        req.params['x-oss-access-key-id'] = credentials.get_access_key_id()
 
-        signature = self.__make_signature(req, bucket_name, key, additional_headers)
+        signature = self.__make_signature(req, bucket_name, key, additional_headers, credentials)
 
         req.params['x-oss-signature'] = signature
 
         return req.url + '?' + '&'.join(_param_to_quoted_query(k, v) for k, v in req.params.items())
 
-    def __make_signature(self, req, bucket_name, key, additional_headers):
+    def __make_signature(self, req, bucket_name, key, additional_headers, credentials):
         if is_py2:
             string_to_sign = self.__get_string_to_sign(req, bucket_name, key, additional_headers)
         else:
             string_to_sign = self.__get_bytes_to_sign(req, bucket_name, key, additional_headers)
 
         logger.debug('Make signature: string to be signed = {0}'.format(string_to_sign))
 
-        h = hmac.new(to_bytes(self.secret), to_bytes(string_to_sign), hashlib.sha256)
+        h = hmac.new(to_bytes(credentials.get_access_key_secret()), to_bytes(string_to_sign), hashlib.sha256)
         return utils.b64encode_as_string(h.digest())
 
     def __get_additional_headers(self, req, in_additional_headers):
@@ -427,7 +455,7 @@ def __get_bytes_to_sign(self, req, bucket_name, key, additional_header_list):
             canonicalized_oss_headers +\
             additional_headers + b'\n' +\
             canonicalized_resource
-    
+
     def __get_canonicalized_oss_headers_bytes(self, req, additional_headers):
         """
         :param additional_headers: 小写的headers列表, 并且这些headers都不以'x-oss-'为前缀.
@@ -442,3 +470,13 @@ def __get_canonicalized_oss_headers_bytes(self, req, additional_headers):
         canon_headers.sort(key=lambda x: x[0])
 
         return b''.join(to_bytes(v[0]) + b':' + to_bytes(v[1]) + b'\n' for v in canon_headers)
+
+
+class AuthV2(ProviderAuthV2):
+    """签名版本2，与版本1的区别在：
+    1. 使用SHA256算法，具有更高的安全性
+    2. 参数计算包含所有的HTTP查询参数
+    """
+    def __init__(self, access_key_id, access_key_secret):
+        credentials_provider = StaticCredentialsProvider(access_key_id.strip(), access_key_secret.strip())
+        super(AuthV2, self).__init__(credentials_provider)
diff --git a/oss2/credentials.py b/oss2/credentials.py
@@ -0,0 +1,132 @@
+# -*- coding: utf-8 -*-
+
+import time
+import requests
+import json
+import logging
+import threading
+from .exceptions import ClientError
+from .utils import to_unixtime
+from .compat import to_unicode
+
+logger = logging.getLogger(__name__)
+
+
+class Credentials(object):
+    def __init__(self, access_key_id="", access_key_secret="", security_token=""):
+        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
+
+
+DEFAULT_ECS_SESSION_TOKEN_DURATION_SECONDS = 3600 * 6
+DEFAULT_ECS_SESSION_EXPIRED_FACTOR = 0.85
+
+
+class EcsRamRoleCredential(Credentials):
+    def __init__(self,
+                 access_key_id,
+                 access_key_secret,
+                 security_token,
+                 expiration,
+                 duration,
+                 expired_factor=None):
+        self.access_key_id = access_key_id
+        self.access_key_secret = access_key_secret
+        self.security_token = security_token
+        self.expiration = expiration
+        self.duration = duration
+        self.expired_factor = expired_factor or DEFAULT_ECS_SESSION_EXPIRED_FACTOR
+
+    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
+
+    def will_soon_expire(self):
+        now = int(time.time())
+        return self.duration * (1.0 - self.expired_factor) > self.expiration - now
+
+
+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
+
+
+class EcsRamRoleCredentialsProvider(CredentialsProvider):
+    def __init__(self, auth_host, max_retries=3, timeout=10):
+        self.fetcher = EcsRamRoleCredentialsFetcher(auth_host)
+        self.max_retries = max_retries
+        self.timeout = timeout
+        self.credentials = None
+        self.__lock = threading.Lock()
+
+    def get_credentials(self):
+        if self.credentials is None or self.credentials.will_soon_expire():
+            with self.__lock:
+                if self.credentials is None or self.credentials.will_soon_expire():
+                    try:
+                        self.credentials = self.fetcher.fetch(self.max_retries, self.timeout)
+                    except Exception as e:
+                        logger.error("Exception: {0}".format(e))
+                        if self.credentials is None:
+                            raise
+
+        return self.credentials
+
+
+class EcsRamRoleCredentialsFetcher(object):
+    def __init__(self, auth_host):
+        self.auth_host = auth_host
+
+    def fetch(self, retry_times=3, timeout=10):
+        for i in range(0, retry_times):
+            try:
+                response = requests.get(self.auth_host, timeout=timeout)
+                if response.status_code != 200:
+                    raise ClientError(
+                        "Failed to fetch credentials url, http code:{0}, msg:{1}".format(response.status_code,
+                                                                                         response.text))
+                dic = json.loads(to_unicode(response.content))
+                code = dic.get('Code')
+                access_key_id = dic.get('AccessKeyId')
+                access_key_secret = dic.get('AccessKeySecret')
+                security_token = dic.get('SecurityToken')
+                expiration_date = dic.get('Expiration')
+                last_updated_date = dic.get('LastUpdated')
+
+                if code != "Success":
+                    raise ClientError("Get credentials from ECS metadata service error, code: {0}".format(code))
+
+                expiration_stamp = to_unixtime(expiration_date, "%Y-%m-%dT%H:%M:%SZ")
+                duration = DEFAULT_ECS_SESSION_TOKEN_DURATION_SECONDS
+                if last_updated_date is not None:
+                    last_updated_stamp = to_unixtime(last_updated_date, "%Y-%m-%dT%H:%M:%SZ")
+                    duration = expiration_stamp - last_updated_stamp
+                return EcsRamRoleCredential(access_key_id, access_key_secret, security_token, expiration_stamp,
+                                            duration, DEFAULT_ECS_SESSION_EXPIRED_FACTOR)
+            except Exception as e:
+                if i == retry_times - 1:
+                    logger.error("Exception: {0}".format(e))
+                    raise ClientError("Failed to get credentials from ECS metadata service. {0}".format(e))
diff --git a/tests/common.py b/tests/common.py
@@ -36,6 +36,7 @@
 OSS_INVENTORY_BUCKET_DESTINATION_ACCOUNT = os.getenv("OSS_TEST_RAM_UID")
 
 OSS_AUTH_VERSION = None
+OSS_TEST_AUTH_SERVER_HOST = os.getenv("OSS_TEST_AUTH_SERVER_HOST")
 
 private_key = RSA.generate(1024)
 public_key = private_key.publickey()
diff --git a/tests/test_credentials_provider.py b/tests/test_credentials_provider.py
diff --git a/tests/test_sts.py b/tests/test_sts.py
