etcd/main.go

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	ini "gopkg.in/ini.v1"
	"io/ioutil"
	"os"
	"reflect"

	"etcd/common.in/clinit"
	"etcd/common.in/config"
	client "go.etcd.io/etcd/client"
	"strings"
)

var (
	configFile = flag.String("config", "./config", "config file location")
	set string
)

func PathExists(path string) bool {
	_, err := os.Stat(path)
	if err == nil {
		return true
	}
	return false
}

func Usage() {
	fmt.Println("etcd:")
	fmt.Println("subcommand:")
	fmt.Println("    set-from-file  filename")
	fmt.Println("    set  key   value")
	fmt.Println("    get  [key]")
	os.Exit(1)
}

func setKey(key, value, encryptKey string) {
	fmt.Println(key,"  ",value)
	if set == "no"{
		return
	}
	
	keyApi := client.NewKeysAPI(clinit.GetEtcdClient())
	opts := &client.SetOptions{PrevExist: client.PrevIgnore}
	if encryptKey != "" {
		aesString, err := config.AesEncrypt(value, encryptKey)
		if err != nil {
			fmt.Printf("AesEncrypt failed. error:%s", err)
			os.Exit(1)
		}

		baseString := config.Base64UrlSafeEncode([]byte(aesString))
		resp, err := keyApi.Set(context.Background(), key, baseString, opts)
		if err != nil {
			fmt.Printf("set etcd failed. response :%s ,error:%s", resp, err)
			os.Exit(1)
		}
	} else {
		resp, err := keyApi.Set(context.Background(), key, value, opts)
		if err != nil {
			fmt.Printf("set etcd failed. response :%s ,error:%s", resp, err)
			os.Exit(1)
		}
	}
}

func SetLine(strList []string, encryptKey string) {
	for _, value := range strList {
		if len(value) == 0 || value[0] == '#' || len(strings.Fields(value)) != 2 {
			continue
		}
		str := strings.Fields(value)
		setKey(str[0], str[1], encryptKey)
	}
}

func setData(basePath string,confData map[string] interface{},encryptKey string){
	for k,v := range confData{
		key := fmt.Sprintf("%s/%s",basePath,k)

		if reflect.ValueOf(v).Kind() == reflect.Map{
			setData(key,v.(map[string]interface{}),encryptKey)
		}else if reflect.ValueOf(v).Kind() == reflect.String{
			setKey(key, v.(string), encryptKey)
		}else{
			fmt.Println(key,":is not map or string")
		}
	}
}

func SetLineNew(basePath string,data []byte, encryptKey string) {
	confData:= make(map[string] interface{})
	e:= json.Unmarshal(data,&confData)

	if e!= nil {
		panic(e)
	}
	setData(basePath,confData,encryptKey)
}

func getNodeData(key string, head *client.Node) (value interface{}) {
	s0 := strings.Split(head.Key, "/")
	len0 := len(s0)
	if len0 == 0 {
		return
	}

	if head.Dir {
		mapData := map[string]interface{}{}
		for _, node := range head.Nodes {
			s1 := strings.Split(node.Key, "/")
			len1 := len(s1)
			if len1 == 0 {
				break
			}

			mapData[s1[len1-1]] = getNodeData(key, node)
		}

		value = mapData
	} else {

		if key != "" && head.Value != "" {
			if bytesData, err := config.Base64URLDecode(head.Value); err != nil {
				fmt.Printf("Base64URLDecode(%s) failed. error:%s", head.Value, err)
				os.Exit(1)
			} else {
				if data, err := config.AesDecrypt(bytesData, []byte(key)); err != nil {
					fmt.Printf("AesDecrypt failed. error:%s", err)
					os.Exit(1)
				} else {
					value = string(data)
				}
			}
		} else {
			value = head.Value
		}
		fmt.Println(head.Key, value)
	}
	return

}

func GetConfig(runmode, key string, cli client.Client) {
	keysAPI := client.NewKeysAPI(cli)
	basePath := fmt.Sprintf("/%s/config", runmode)
	if resp, err := keysAPI.Get(context.Background(), basePath, &client.GetOptions{
		Recursive: true,
	}); err == nil && resp != nil && resp.Node != nil {
		getNodeData(key, resp.Node)
	} else {
		fmt.Printf("get %s failed. error:%s", basePath, err)
		os.Exit(1)
	}

	return
}

func main() {
	flag.Parse()

	cfg, err := ini.Load("./app.conf")
	if err != nil {
		fmt.Printf("Fail to read file: %v\n\n", err)
		os.Exit(1)
	}

	etcdAddrs := strings.Split(cfg.Section("").Key("etcd_addrs").String(), ",")
	runmode := cfg.Section("").Key("runmode").String()
	encryptKey := cfg.Section("").Key("encrypt_key").String()
	projectName := cfg.Section("").Key("project_name").String()
	set = cfg.Section("").Key("set").String()
	clinit.InitEtcd(etcdAddrs)
	if len(os.Args) >= 2 {
		if os.Args[1] == "set-from-file" && len(os.Args) == 3 {
			isExist := PathExists(os.Args[2])
			if !isExist {
				Usage()
				os.Exit(1)
			}
			b, err := ioutil.ReadFile(os.Args[2])
			if err != nil {
				fmt.Print(err)
				os.Exit(1)
			}
			basePath := fmt.Sprintf("/%s/%s/config",projectName,runmode)
			SetLineNew(basePath,b, encryptKey)
			//SetLine(strings.Split(string(b), "\n"), encryptKey)
		} else if os.Args[1] == "set" {
			if len(os.Args) != 4 {
				Usage()
				os.Exit(1)
			}

			setKey(os.Args[2], os.Args[3], encryptKey)

		} else if os.Args[1] == "get" {
			opts := &client.GetOptions{Recursive: true}

			if len(os.Args) == 3 {
				keyApi := client.NewKeysAPI(clinit.GetEtcdClient())
				resp, err := keyApi.Get(context.Background(), os.Args[2], opts)
				if err == nil && resp != nil && resp.Node != nil {
					getNodeData(encryptKey, resp.Node)
				} else {
					fmt.Printf("get %s failed. error:%s", os.Args[2], err)
					os.Exit(1)
				}
			} else if len(os.Args) == 2 {
				GetConfig(runmode, encryptKey, clinit.GetEtcdClient())
			} else {
				Usage()
				os.Exit(1)
			}

		} else {
			fmt.Println("not support subcommand: " + os.Args[1])
			Usage()
		}
	} else {
		Usage()
	}
}