key-manager/main.go

package main

import (
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/console/prompt"
	tronaddr "github.com/fbsobreira/gotron-sdk/pkg/address"
	logging "github.com/ipfs/go-log/v2"
	"github.com/urfave/cli/v2"
	"google.golang.org/grpc"
	"gorm.io/gorm"
	"io/ioutil"
	"key-manager/conf"
	"key-manager/crypto"
	"key-manager/dao"
	"key-manager/service"
	"net"
	"os"
	"strings"
)

var log = logging.Logger("main")

func main() {
	_ = logging.SetLogLevel("*", "INFO")

	app := cli.App{
		Name:  "key-manager",
		Usage: "eth and tron key manager service",
		Commands: []*cli.Command{
			initCmd,
			runCmd,
			whitelistCmd,
		},
		EnableBashCompletion: true,
	}

	if err := app.Run(os.Args); err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n\n", err)
		os.Exit(1)
	}
}

var initCmd = &cli.Command{
	Name:  "init",
	Usage: "init key-manager service",
	Flags: []cli.Flag{
		&cli.StringFlag{
			Name:  "config",
			Usage: "config file path",
		},
	},
	Action: func(cctx *cli.Context) error {
		err := conf.InitConfig(cctx.String("config"))
		if err != nil {
			return err
		}
		db, err := dao.InitMysqlDB()
		if err != nil {
			return err
		}

		_, err = db.GetPassword()
		if err == nil {
			return errors.New("password already set")
		}

		if !errors.Is(err, gorm.ErrRecordNotFound) {
			return err
		}

		log.Info("Please set a password for the encrypted mnemonic")

		password, err := stdinPassword(true)
		if err != nil {
			return err
		}
		scrypt := crypto.Scrypt(password)
		scryptKey := hex.EncodeToString(scrypt)
		err = db.CreatePassword(&dao.Password{
			Password: scryptKey,
		})

		if err != nil {
			return err
		}

		return nil
	},
}

var runCmd = &cli.Command{
	Name:  "run",
	Usage: "run key-manager process",
	Flags: []cli.Flag{
		&cli.StringFlag{
			Name:  "listen",
			Usage: "The host address and port on which the key manager will listen",
			Value: "127.0.0.1:5556",
		},
		&cli.StringFlag{
			Name:  "config",
			Usage: "config file path",
		},
		&cli.StringFlag{
			Name:  "password",
			Usage: "password file path",
		},
	},
	Action: func(cctx *cli.Context) error {
		err := conf.InitConfig(cctx.String("config"))
		if err != nil {
			return err
		}
		db, err := dao.InitMysqlDB()
		if err != nil {
			return err
		}

		scrypt, err := db.GetPassword()
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return errors.New("please set a password first. run: key-manager init")
		} else if err != nil {
			return err
		}

		scryptKey, err := hex.DecodeString(scrypt)
		if err != nil {
			return err
		}

		var password string

		if cctx.IsSet("password") {
			p := cctx.String("password")
			data, err := ioutil.ReadFile(p)
			if err != nil {
				return err
			}
			password = strings.TrimSpace(string(data))
		} else {
			var err error
			password, err = stdinPassword(false)
			if err != nil {
				return err
			}
		}

		ok, err := crypto.VerifyScrypt(password, scryptKey)
		if err != nil {
			return err
		}
		if !ok {
			return errors.New("password verification failed")
		}

		listen := cctx.String("listen")
		listener, err := net.Listen("tcp", listen)
		if err != nil {
			return err
		}

		log.Infof("grpc server Listing on: %s", listen)

		grpcServer := grpc.NewServer()
		server, err := service.NewKeyManager(db, password)
		if err != nil {
			return err
		}

		service.RegisterKeyManagerServer(grpcServer, server)
		if err = grpcServer.Serve(listener); err != nil {
			log.Error(err)
			return err
		}

		return nil
	},
}

var whitelistCmd = &cli.Command{
	Name:  "whitelist",
	Usage: "whitelist tools",
	Subcommands: []*cli.Command{
		listCmd,
		addCmd,
		deleteCmd,
	},
}

var addCmd = &cli.Command{
	Name:      "add",
	Usage:     "add whitelist address",
	ArgsUsage: "[name] [address]",
	Flags: []cli.Flag{
		&cli.StringFlag{
			Name:  "config",
			Usage: "config file path",
		},
	},
	Action: func(cctx *cli.Context) error {
		err := conf.InitConfig(cctx.String("config"))
		if err != nil {
			return err
		}
		db, err := dao.InitMysqlDB()
		if err != nil {
			return err
		}
		if cctx.NArg() != 2 {
			return errors.New("params need [name] [address]")
		}

		name := cctx.Args().Get(0)
		address := cctx.Args().Get(1)
		if address[:2] == "0x" {
			address = common.HexToAddress(address).String()
		} else {
			addr, err := tronaddr.Base58ToAddress(address)
			if err != nil {
				return err
			}
			address = addr.String()
		}

		err = db.CreateWhitelist(&dao.Whitelist{
			Name:         name,
			WhiteAddress: address,
		})
		if err != nil {
			return err
		}

		fmt.Printf("add whitelist address success, name: %s, addr: %s \n", name, address)

		return nil
	},
}

var deleteCmd = &cli.Command{
	Name:      "delete",
	Usage:     "delete whitelist address",
	ArgsUsage: "[name]",
	Flags: []cli.Flag{
		&cli.StringFlag{
			Name:  "config",
			Usage: "config file path",
		},
	},
	Action: func(cctx *cli.Context) error {
		err := conf.InitConfig(cctx.String("config"))
		if err != nil {
			return err
		}
		db, err := dao.InitMysqlDB()
		if err != nil {
			return err
		}
		if cctx.NArg() != 1 {
			return errors.New("params need [name]")
		}

		name := cctx.Args().Get(0)

		err = db.DeleteWhitelist(name)
		if err != nil {
			return err
		}

		fmt.Printf("delete whitelist address success, name: %s \n", name)

		return nil
	},
}

var listCmd = &cli.Command{
	Name:  "list",
	Usage: "list whitelist address",
	Flags: []cli.Flag{
		&cli.StringFlag{
			Name:  "config",
			Usage: "config file path",
		},
	},
	Action: func(cctx *cli.Context) error {
		err := conf.InitConfig(cctx.String("config"))
		if err != nil {
			return err
		}
		db, err := dao.InitMysqlDB()
		if err != nil {
			return err
		}

		list, err := db.GetAllWhitelists()
		if err != nil {
			return err
		}

		for _, l := range list {
			fmt.Printf("whitelist address, name: %s, addr: %s \n", l.Name, l.WhiteAddress)
		}

		return nil
	},
}

func stdinPassword(isConfirm bool) (string, error) {
	password, err := prompt.Stdin.PromptPassword("Password: ")
	if err != nil {
		return "", err
	}

	if isConfirm {
		confirm, err := prompt.Stdin.PromptPassword("Confirm Password: ")
		if err != nil {
			return "", fmt.Errorf("failed to read password confirmation: %v", err)
		}
		if password != confirm {
			return "", errors.New("passwords do not match")
		}
	}

	return password, nil
}