[add] token update util

util/token.go

@@ -0,0 +1,190 @@
+package util
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"encoding/base64"
+	"encoding/json"
+	"errors"
+	"io/ioutil"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+
+	sls "github.com/aliyun/aliyun-log-go-sdk"
+)
+
+const (
+	aliyunECSRamURL      = "http://100.100.100.200/latest/meta-data/ram/security-credentials/"
+	expirationTimeFormat = "2006-01-02T15:04:05Z"
+)
+
+var errNoFile = errors.New("no secret file")
+
+// AKInfo ...
+type AKInfo struct {
+	AccessKeyId     string `json:"access.key.id"`
+	AccessKeySecret string `json:"access.key.secret"`
+	SecurityToken   string `json:"security.token"`
+	Expiration      string `json:"expiration"`
+	Keyring         string `json:"keyring"`
+}
+
+// SecurityTokenResult ...
+type SecurityTokenResult struct {
+	AccessKeyId     string
+	AccessKeySecret string
+	Expiration      string
+	SecurityToken   string
+	Code            string
+	LastUpdated     string
+}
+
+func getToken() (result []byte, err error) {
+	client := http.Client{
+		Timeout: time.Second * 3,
+	}
+	var respList *http.Response
+	respList, err = client.Get(aliyunECSRamURL)
+	if err != nil {
+		return nil, err
+	}
+	defer respList.Body.Close()
+	var body []byte
+	body, err = ioutil.ReadAll(respList.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	bodyStr := string(body)
+	bodyStr = strings.TrimSpace(bodyStr)
+	roles := strings.Split(bodyStr, "\n")
+	role := roles[0]
+
+	var respGet *http.Response
+	respGet, err = client.Get(aliyunECSRamURL + role)
+	if err != nil {
+		return nil, err
+	}
+	defer respGet.Body.Close()
+	body, err = ioutil.ReadAll(respGet.Body)
+	if err != nil {
+		return nil, err
+	}
+	return body, nil
+}
+func pkcs5UnPadding(origData []byte) []byte {
+	length := len(origData)
+	unpadding := int(origData[length-1])
+	return origData[:(length - unpadding)]
+}
+
+func decrypt(s string, keyring []byte) ([]byte, error) {
+	cdata, err := base64.StdEncoding.DecodeString(s)
+	if err != nil {
+		return nil, err
+	}
+	block, err := aes.NewCipher(keyring)
+	if err != nil {
+		return nil, err
+	}
+	blockSize := block.BlockSize()
+
+	iv := cdata[:blockSize]
+	blockMode := cipher.NewCBCDecrypter(block, iv)
+	origData := make([]byte, len(cdata)-blockSize)
+
+	blockMode.CryptBlocks(origData, cdata[blockSize:])
+
+	origData = pkcs5UnPadding(origData)
+	return origData, nil
+}
+
+func getAKFromLocalFile(configFilePath string) (accessKeyID, accessKeySecret, securityToken string, expireTime time.Time, err error) {
+	if _, err = os.Stat(configFilePath); err == nil {
+		var akInfo AKInfo
+		//获取token config json
+		encodeTokenCfg, err := ioutil.ReadFile(configFilePath)
+		if err != nil {
+			return accessKeyID, accessKeySecret, securityToken, expireTime, err
+		}
+		err = json.Unmarshal(encodeTokenCfg, &akInfo)
+		if err != nil {
+			return accessKeyID, accessKeySecret, securityToken, expireTime, err
+		}
+		keyring := akInfo.Keyring
+		ak, err := decrypt(akInfo.AccessKeyId, []byte(keyring))
+		if err != nil {
+			return accessKeyID, accessKeySecret, securityToken, expireTime, err
+		}
+
+		sk, err := decrypt(akInfo.AccessKeySecret, []byte(keyring))
+		if err != nil {
+			return accessKeyID, accessKeySecret, securityToken, expireTime, err
+		}
+
+		token, err := decrypt(akInfo.SecurityToken, []byte(keyring))
+		if err != nil {
+			return accessKeyID, accessKeySecret, securityToken, expireTime, err
+		}
+		layout := "2006-01-02T15:04:05Z"
+		t, err := time.Parse(layout, akInfo.Expiration)
+		if err != nil {
+			return accessKeyID, accessKeySecret, securityToken, expireTime, err
+		}
+		if t.Before(time.Now()) {
+			err = errors.New("invalid token which is expired")
+		}
+		akInfo.AccessKeyId = string(ak)
+		akInfo.AccessKeySecret = string(sk)
+		akInfo.SecurityToken = string(token)
+
+		if err != nil {
+			return accessKeyID, accessKeySecret, securityToken, expireTime, err
+		}
+		return akInfo.AccessKeyId, akInfo.AccessKeySecret, akInfo.SecurityToken, t, nil
+	}
+	return accessKeyID, accessKeySecret, securityToken, expireTime, errNoFile
+}
+
+func updateTokenFunction(configFilePath string) (accessKeyID, accessKeySecret, securityToken string, expireTime time.Time, err error) {
+	if configFilePath != "" {
+		accessKeyID, accessKeySecret, securityToken, expireTime, err = getAKFromLocalFile(configFilePath)
+		if err != errNoFile {
+			return accessKeyID, accessKeySecret, securityToken, expireTime, err
+		}
+	}
+	var tokenResultBuffer []byte
+	for tryTime := 0; tryTime < 3; tryTime++ {
+		tokenResultBuffer, err = getToken()
+		if err != nil {
+			continue
+		}
+		var tokenResult SecurityTokenResult
+		err = json.Unmarshal(tokenResultBuffer, &tokenResult)
+		if err != nil {
+			continue
+		}
+		if strings.ToLower(tokenResult.Code) != "success" {
+			tokenResult.AccessKeySecret = "x"
+			tokenResult.SecurityToken = "x"
+			continue
+		}
+		expireTime, err := time.Parse(expirationTimeFormat, tokenResult.Expiration)
+		if err != nil {
+			tokenResult.AccessKeySecret = "x"
+			tokenResult.SecurityToken = "x"
+			continue
+		}
+		return tokenResult.AccessKeyId, tokenResult.AccessKeySecret, tokenResult.SecurityToken, expireTime, nil
+	}
+	return accessKeyID, accessKeySecret, securityToken, expireTime, err
+}
+
+// NewTokenUpdateFunc create a token update function for ACK or ECS
+func NewTokenUpdateFunc(role string, configFilePath string) (tokenUpdateFunc sls.UpdateTokenFunction, shutdown <-chan struct{}) {
+	return func() (accessKeyID string, accessKeySecret string, securityToken string, expireTime time.Time, err error) {
+		return updateTokenFunction(configFilePath)
+	}, make(<-chan struct{})
+}