smart-government-management/utils/utils.go

// Copyright 2019  github.com. All rights reserved.
// Use of this source code is governed by  github.com.
package utils

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/md5"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"github.com/dgrijalva/jwt-go"
	"github.com/tidwall/gjson"
	"math/rand"
	"strings"
	"time"
	"github.com/jaryhe/gopkgs/cache"
)

var jwtSecret = []byte("12345678abcdefg")
//var CRYPTO_KEY = string(RandStr(32, 3))

func Base64URLDecode(data string) ([]byte, error) {
	var missing = (4 - len(data)%4) % 4
	data += strings.Repeat("=", missing)
	return base64.URLEncoding.DecodeString(data)
}

func Base64UrlSafeEncode(source []byte) string {
	// Base64 Url Safe is the same as Base64 but does not contain '/' and '+' (replaced by '_' and '-') and trailing '=' are removed.
	bytearr := base64.StdEncoding.EncodeToString(source)
	safeurl := strings.Replace(string(bytearr), "/", "_", -1)
	safeurl = strings.Replace(safeurl, "+", "-", -1)
	safeurl = strings.Replace(safeurl, "=", "", -1)
	return safeurl
}

func AesDecrypt(crypted, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	blockMode := NewECBDecrypter(block)
	origData := make([]byte, len(crypted))
	blockMode.CryptBlocks(origData, crypted)
	origData = PKCS5UnPadding(origData)
	return origData, nil
}

func AesEncrypt(src, key string) ([]byte, error) {
	block, err := aes.NewCipher([]byte(key))
	if err != nil {
		return nil, err
	}

	ecb := NewECBEncrypter(block)
	content := []byte(src)
	content = PKCS5Padding(content, block.BlockSize())
	crypted := make([]byte, len(content))
	ecb.CryptBlocks(crypted, content)

	return crypted, nil
}

func PKCS5Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

func PKCS5UnPadding(origData []byte) []byte {
	length := len(origData)
	unpadding := int(origData[length-1])
	return origData[:(length - unpadding)]
}

type ecb struct {
	b         cipher.Block
	blockSize int
}

func newECB(b cipher.Block) *ecb {
	return &ecb{
		b:         b,
		blockSize: b.BlockSize(),
	}
}

type ecbEncrypter ecb

// NewECBEncrypter returns a BlockMode which encrypts in electronic code book mode, using the given Block.
func NewECBEncrypter(b cipher.Block) cipher.BlockMode {
	return (*ecbEncrypter)(newECB(b))
}

func (x *ecbEncrypter) BlockSize() int { return x.blockSize }

func (x *ecbEncrypter) CryptBlocks(dst, src []byte) {
	if len(src)%x.blockSize != 0 {
		fmt.Println("crypto/cipher: input not full blocks")
		return
	}
	if len(dst) < len(src) {
		fmt.Println("crypto/cipher: output smaller than input")
		return
	}
	for len(src) > 0 {
		x.b.Encrypt(dst, src[:x.blockSize])
		src = src[x.blockSize:]
		dst = dst[x.blockSize:]
	}
}

type ecbDecrypter ecb

// NewECBDecrypter returns a BlockMode which decrypts in electronic code book mode, using the given Block.
func NewECBDecrypter(b cipher.Block) cipher.BlockMode {
	return (*ecbDecrypter)(newECB(b))
}

func (x *ecbDecrypter) BlockSize() int { return x.blockSize }

func (x *ecbDecrypter) CryptBlocks(dst, src []byte) {
	if len(src)%x.blockSize != 0 {
		fmt.Println("crypto/cipher: input not full blocks")
		return
	}
	if len(dst) < len(src) {
		fmt.Println("crypto/cipher: output smaller than input")
		return
	}
	for len(src) > 0 {
		x.b.Decrypt(dst, src[:x.blockSize])
		src = src[x.blockSize:]
		dst = dst[x.blockSize:]
	}
}

func RandStr(size int, kind int) []byte {
	ikind, kinds, result := kind, [][]int{[]int{10, 48}, []int{26, 97}, []int{26, 65}}, make([]byte, size)
	is_all := kind > 2 || kind < 0
	rand.Seed(time.Now().UnixNano())
	for i := 0; i < size; i++ {
		if is_all { // random ikind
			ikind = rand.Intn(3)
		}
		scope, base := kinds[ikind][0], kinds[ikind][1]
		result[i] = uint8(base + rand.Intn(scope))
	}
	return result
}

func GenToken(id, issuer, subject string, seconds time.Duration) (string, error) {
	if len(jwtSecret) == 0 {
		return "", fmt.Errorf("jwtSecret is empty.")
	}

	nowTime := time.Now()
	expireTime := nowTime.Add(seconds)

	claims := jwt.StandardClaims{
		ExpiresAt: expireTime.Unix(),
		Id:        id,
		Issuer:    issuer,
		Subject:   subject,
	}

	tokenClaims := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	token, err := tokenClaims.SignedString(jwtSecret)

	return token, err
}

func ParseToken(token string) (*jwt.StandardClaims, error) {
	tokenClaims, err := jwt.ParseWithClaims(
		token,
		&jwt.StandardClaims{},
		func(token *jwt.Token) (interface{}, error) {
			return jwtSecret, nil
		},
	)

	if tokenClaims != nil {
		claims, ok := tokenClaims.Claims.(*jwt.StandardClaims)
		if !ok {
			return nil, err
		}

		if tokenClaims.Valid {
			return claims, nil
		}

		return claims, err
	}

	return nil, err
}

func MD5(s string) string {
	m := md5.New()
	m.Write([]byte(s))
	return hex.EncodeToString(m.Sum(nil))
}

func getMd5Pass(passwd string, crykey string) (string) {
	cryPasswd, err := AesEncrypt(passwd, crykey)
	if err != nil {
		return ""
	}
	ret := MD5(base64.StdEncoding.EncodeToString(cryPasswd))
	return ret
}

func GetCredential() string {
	crykey := string(RandStr(32, 3))
	key := string(RandStr(8, 3))
	md5Pass := getMd5Pass(key, crykey)
	sub := map[string]string{
		"key":key,
		"crykey":crykey,
		"pass":md5Pass,
	}

	subStr, _ := json.Marshal(sub)
	token, err := GenToken("1", "", string(subStr), 600 *time.Second)
	fmt.Printf("%v\n", err)
	return token
}


var credentialKey = "credential_"

func checkCredentialToCache(token string) (bool, error) {
	key := credentialKey + token
	res, err := cache.Redis().Get(key)
	if err != nil && strings.Contains(err.Error(), "redis: nil") == false {
		return false, err
	}
	if res == "1" {
		return true, nil
	}
	return false, nil
}

func saveCredentialToCache(token string) error {
	key := credentialKey + token
	_, err := cache.Redis().SetEx(key, 60*20, "1")
	if err != nil {
		return err
	}
	return nil
}

func VerifyCredential(token string) bool {
	if token == "" {
		return false
	}
	exist, err := checkCredentialToCache(token)
	if err != nil {
		fmt.Printf("%v\n", err)
		return false
	}
	if exist {
		fmt.Printf("凭据已失效\n")
	}
	claimes, err := ParseToken(token)
	if err != nil {
		fmt.Printf("%v\n", err)
		return false
	}
	key := gjson.GetBytes([]byte(claimes.Subject), "key").String()
	crykey := gjson.GetBytes([]byte(claimes.Subject), "crykey").String()
	pass := gjson.GetBytes([]byte(claimes.Subject), "pass").String()

	dst := getMd5Pass(key, crykey)
	if pass != dst || dst == "" {
		fmt.Printf("key is not match\n")
		return false
	}
	if saveCredentialToCache(token) != nil {
		fmt.Printf("内部错误\n")
		return false
	}
	return true
}