jypkg/jyutil/sessionEncryption.go

package jyutil

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"jy/src/jfw/config"
	"jy/src/jfw/nodemgr"
	"log"
	"sync"
	"time"

	util "app.yhyue.com/moapp/jybase/common"
	. "app.yhyue.com/moapp/jybase/date"
	"app.yhyue.com/moapp/jybase/redis"
	"github.com/robfig/cron/v3"
)

//
type RsaStruct struct {
	Sign       string
	PrivateKey *rsa.PrivateKey
	PublicKey  *rsa.PublicKey
	RsaConfig  *RsaConfig
}

type RsaConfig struct {
	RsaSignflag   string
	PrivateKeyStr string
	PublicKeyStr  string
}

type MyWrite struct {
	Byte *bytes.Buffer
}

func (m MyWrite) Write(p []byte) (n int, err error) {
	n, err = m.Byte.Write(p)
	return
}

func getNewRsaConfig() *RsaConfig {
	//私钥-----
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		log.Println(err)
	}
	X509PrivateKey := x509.MarshalPKCS1PrivateKey(privateKey)
	privateBlock := pem.Block{Type: "RSA Private Key", Bytes: X509PrivateKey}
	privateWrite := MyWrite{Byte: bytes.NewBuffer([]byte(""))}
	pem.Encode(privateWrite, &privateBlock)
	privateKeyByte := privateWrite.Byte.Bytes()

	//公钥----
	publicKey := privateKey.PublicKey
	X509PublicKey, err := x509.MarshalPKIXPublicKey(&publicKey)
	if err != nil {
		log.Println(err)
	}
	publicBlock := pem.Block{Type: "RSA Public Key", Bytes: X509PublicKey}
	publicWrite := MyWrite{Byte: bytes.NewBuffer([]byte(""))}
	pem.Encode(publicWrite, &publicBlock)
	publicKeyByte := publicWrite.Byte.Bytes()

	return &RsaConfig{
		RsaSignflag:   time.Now().Format(Date_Full_Layout),
		PrivateKeyStr: string(privateKeyByte),
		PublicKeyStr:  string(publicKeyByte),
	}
}

func initRsa(config *RsaConfig) *RsaStruct {
	//初始化私钥
	private_block, _ := pem.Decode([]byte(config.PrivateKeyStr))
	private, err := x509.ParsePKCS1PrivateKey(private_block.Bytes)
	if err != nil {
		log.Println(err)
	}
	//初始化公钥
	public_block, _ := pem.Decode([]byte(config.PublicKeyStr))
	publicKeyInterface, err := x509.ParsePKIXPublicKey(public_block.Bytes)
	if err != nil {
		log.Println(err)
	}
	public := publicKeyInterface.(*rsa.PublicKey)
	return &RsaStruct{
		Sign:       config.RsaSignflag,
		RsaConfig:  config,
		PrivateKey: private,
		PublicKey:  public,
	}
}

//加密
func (r *RsaStruct) Encrypt(str string) string {
	enStrByte, err := rsa.EncryptPKCS1v15(rand.Reader, r.PublicKey, []byte(str))
	if err != nil {
		log.Printf("%s加密失败%v", str, err)
	}
	return string(enStrByte)
}

//解密
func (r *RsaStruct) Decrypt(enStr string) string {
	deStrByte, err := rsa.DecryptPKCS1v15(rand.Reader, r.PrivateKey, []byte(enStr))
	if err != nil {
		log.Printf("%s解密失败%v", enStr, err)
	}
	return string(deStrByte)
}

type EncryptionStruct struct {
	RsaStruct_Now  *RsaStruct
	MapRsa         map[string]*RsaStruct
	EncryptionConf *EncryptionConf
	Lock           *sync.Mutex
}

type EncryptionConf struct {
	Now        string
	MapRsaConf map[string]*RsaConfig
}

//首次启动加载秘钥
func InitEncryptionStruct() *EncryptionStruct {
	encryStruct := EncryptionStruct{EncryptionConf: &EncryptionConf{}, Lock: &sync.Mutex{}}
	//	log.Println("xxxxxx", redis.GetInt("other", "jynode_"+nodemgr.NodeName))
	//	isMaster := util.If(redis.GetInt("other", "jynode_"+nodemgr.NodeName) == 1, true, false).(bool)
	go func() {
		cn := cron.New()
		str := util.ObjToString(config.Sysconfig["sessionEnctryFlushCorn"])
		cn.AddFunc(str, func() {
			encryStruct.Flush()
		})
		cn.Start()
	}()

	//for {
	redis_obj := util.ObjToString(redis.Get("session", "sessionEncryption"))
	if redis_obj != "" {
		encryconf := EncryptionConf{}
		err := json.Unmarshal([]byte(redis_obj), &encryconf)
		if err != nil {
			log.Println("InitEncryptionStruct Unmarshal redis err", err)
		} else {
			encryStruct.EncryptionConf = &encryconf
			encryStruct.initAllRsa()
			log.Println("从redis中加载EncryptionStruct")
			return &encryStruct
		}
	}

	log.Println("主结点首次创建EncryptionStruct")
	newRsaConf := getNewRsaConfig()
	encryStruct.EncryptionConf.Now = newRsaConf.RsaSignflag
	encryStruct.EncryptionConf.MapRsaConf = map[string]*RsaConfig{newRsaConf.RsaSignflag: newRsaConf}
	encryStruct.initAllRsa()
	go func() {
		//存入redis
		byteArr, err := json.Marshal(encryStruct.EncryptionConf)
		log.Println(string(byteArr))
		if err != nil {
			log.Printf("InitEncryptionStruct save %s redis err %:v\n", byteArr, err)
			return
		}
		redis.Put("session", "sessionEncryption", string(byteArr), -1)
	}()
	return &encryStruct
}

func (e *EncryptionStruct) initAllRsa() {
	e.Lock.Lock()
	defer e.Lock.Unlock()
	e.MapRsa = map[string]*RsaStruct{}
	for k, v := range e.EncryptionConf.MapRsaConf {
		thisRsa := initRsa(v)
		e.MapRsa[k] = thisRsa
		if k == e.EncryptionConf.Now {
			e.RsaStruct_Now = thisRsa
		}
	}
}

//定时更新秘钥
func (e *EncryptionStruct) Flush() {
	isMaster := util.If(redis.GetInt("other", "jynode_"+nodemgr.NodeName) == 1, true, false).(bool)
	if !isMaster {
		//从节点从redis中更新秘钥
		for {
			redis_obj := util.ObjToString(redis.Get("session", "sessionEncryption"))
			if redis_obj != "" {
				encryconf := EncryptionConf{}
				err := json.Unmarshal([]byte(redis_obj), &encryconf)
				if e.RsaStruct_Now.Sign == encryconf.Now {
					log.Println("从节点定时更新EncryptionStruct失败，一分钟后重试", e.RsaStruct_Now.Sign)
					time.Sleep(time.Minute)
					if redis.GetInt("other", "jynode_"+nodemgr.NodeName) == 1 {
						break
					} else {
						continue
					}
				}
				if err != nil {
					log.Println("InitEncryptionStruct Unmarshal redis err", err)
				} else {
					e.EncryptionConf = &encryconf
					e.initAllRsa()
					log.Println("从节点定时更新EncryptionStruct")
					return
				}
			}
		}
	} else {
		e.Lock.Lock()
		//主结点生成秘钥 存入redis
		log.Printf("主结点更新EncryptionStruct before  %+v", e.MapRsa)
		last := e.RsaStruct_Now
		newRsa := initRsa(getNewRsaConfig())
		//更新map
		e.MapRsa = map[string]*RsaStruct{newRsa.Sign: newRsa, last.Sign: last}
		//更新默认加密
		e.RsaStruct_Now = newRsa
		log.Printf("after  %+v", e.MapRsa)
		e.Lock.Unlock()
		//更新redis
		func() {
			e.EncryptionConf.Now = e.RsaStruct_Now.Sign
			e.EncryptionConf.MapRsaConf = map[string]*RsaConfig{newRsa.Sign: newRsa.RsaConfig, last.Sign: last.RsaConfig}
			byteArr, err := json.Marshal(e.EncryptionConf)
			log.Println("主结点更新sessionEncryption", string(byteArr))
			if err != nil {
				log.Printf("InitEncryptionStruct save %s redis err %:v\n", byteArr, err)
				return
			}
			redis.Put("session", "sessionEncryption", string(byteArr), -1)
		}()
	}
}