key-manager/service/service.go

package service

import (
	"crypto/ecdsa"
	"encoding/hex"
	"errors"
	logging "github.com/ipfs/go-log/v2"
	"key-manager/crypto"
	"key-manager/dao"
	"key-manager/hd"
	"sync"
)

var log = logging.Logger("key-manager")

type KeyManager struct {
	UnimplementedKeyManagerServer
	lk         sync.Mutex
	dao        *dao.Dao
	password   string
	whitelists map[string]struct{}          // key = address
	mnemonics  map[string]string            // key = name
	signers    map[string]*ecdsa.PrivateKey // key = address
}

func NewKeyManager(dao *dao.Dao, password string) (*KeyManager, error) {
	var km = &KeyManager{
		UnimplementedKeyManagerServer: UnimplementedKeyManagerServer{},
		lk:                            sync.Mutex{},
		dao:                           dao,
		password:                      password,
		whitelists:                    map[string]struct{}{},
		mnemonics:                     map[string]string{},
		signers:                       map[string]*ecdsa.PrivateKey{},
	}

	whitelists, err := dao.GetAllWhitelists()
	if err != nil {
		return nil, err
	}

	for _, whitelist := range whitelists {
		km.whitelists[whitelist.WhiteAddress] = struct{}{}
		log.Warnf("whitelist address: %s - %s", whitelist.Name, whitelist.WhiteAddress)
	}

	encryptedMnemonics, err := dao.GetAllMnemonics()
	if err != nil {
		return nil, err
	}

	passwordKey := crypto.GenerateEncryptKey([]byte(password))
	for _, encryptedMnemonic := range encryptedMnemonics {
		mnemonic, err := decryptMnemonic(encryptedMnemonic.Mnemonic, passwordKey)
		if err != nil {
			log.Errorw("decryptMnemonic failed", "name", encryptedMnemonic.Name, "encryptedMnemonic", encryptedMnemonic.Mnemonic)
			return nil, err
		}
		if ok := hd.CheckMnemonic(mnemonic); !ok {
			return nil, errors.New("invalid mnemonic")
		}
		km.mnemonics[encryptedMnemonic.Name] = mnemonic
		log.Infow("load mnemonic", "name", encryptedMnemonic.Name)
	}

	return km, nil
}

func encryptMnemonic(mnemonic string, key []byte) (string, error) {
	encryptedMnemonic, err := crypto.Encrypt([]byte(mnemonic), key)
	if err != nil {
		return "", err
	}

	return hex.EncodeToString(encryptedMnemonic), err
}

func decryptMnemonic(encryptedData string, key []byte) (string, error) {
	encryptedMnemonic, err := hex.DecodeString(encryptedData)
	if err != nil {
		return "", err
	}
	mnemonic, err := crypto.Decrypt(encryptedMnemonic, key)
	if err != nil {
		return "", err
	}

	return string(mnemonic), nil
}

func checkNetwork(network string) error {
	switch network {
	case eth, tron:
		return nil
	default:
		return errors.New("network must be eth or tron")
	}
}