From ce51ad31b46cb5ab4dc3bc782c57778c000d4497 Mon Sep 17 00:00:00 2001 From: "davidzhang.zc" Date: Mon, 11 Jan 2021 14:37:29 +0800 Subject: [PATCH] [add] token update util --- util/token.go | 190 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 190 insertions(+) create mode 100644 util/token.go diff --git a/util/token.go b/util/token.go new file mode 100644 index 00000000..6c51ffb9 --- /dev/null +++ b/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{}) +}