qcloudcos/main.go

package main

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"github.com/AlecAivazis/survey"
	"github.com/tencentyun/cos-go-sdk-v5"
	"github.com/wailovet/osmanthuswine/src/helper"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"runtime"
	"strings"
)

var config TencentYunConfig

var op string
var path string
var qpath string
var key string

func loadKey() {
	homePath, err := home()
	if err != nil {
		fmt.Println("错误 102: ", err)
		return
	}
	keyFile := filepath.Join(homePath, ".qcloud-key")
	data, _ := ioutil.ReadFile(keyFile)
	if string(data) == "" {
		fmt.Println("错误 204 找不到key,请使用使用set-key进行初始化.")
		return
	}
	key = string(data)
}

func initKey() {

	if key == "" {
		fmt.Println("key is empty!")
		os.Exit(1)
	}
	var err error
	var dekey []byte
	dekey, err = base64.StdEncoding.DecodeString(key)
	if err != nil {
		fmt.Println("错误 201:", err)
		os.Exit(1)
	}

	dekey, err = helper.DefaultAesDecrypt(dekey)
	if err != nil {
		fmt.Println("错误 202:", err)
		os.Exit(1)
	}
	err = json.Unmarshal(dekey, &config)
	if err != nil {
		fmt.Println("错误 203:", err)
		os.Exit(1)
	}
}

func main() {

	helper.AesKey = []byte("3136352472abcdef")
	flag.StringVar(&op, "op", "", "操作:make-key[生成密钥] | set-key[设置密钥] | upload [上传文件] | list [获取文件]")
	flag.StringVar(&path, "path", "", "上传目录")
	flag.StringVar(&qpath, "qpath", "", "云上传目的目录")
	flag.StringVar(&key, "key", "", "key")
	flag.Parse()

	switch op {
	case "set-key":
		key = askText("输入key")
		initKey()
		_, err := getTencentYunCosFiles("")
		if err != nil {
			fmt.Println("key error 101:", err)
			return
		}

		homePath, err := home()
		if err != nil {
			fmt.Println("错误 102:", err)
			return
		}
		keyFile := filepath.Join(homePath, ".qcloud-key")
		err = ioutil.WriteFile(keyFile, []byte(key), 0644)
		if err != nil {
			fmt.Println("错误 103:", err)
			return
		}

		fmt.Println("完成!")
	case "make-key":
		config.SecretID = askText("输入SecretID")
		config.SecretKey = askText("输入SecretKey")
		config.BucketURL = askText("输入BucketURL")
		config.Root = askText("输入Root")

		mkey, err := helper.DefaultAesEncrypt([]byte(helper.JsonEncode(config)))
		if err != nil {
			fmt.Println("错误 104:", err)
			return
		}
		mkey = []byte(base64.StdEncoding.EncodeToString(mkey))
		fmt.Println("key:", string(mkey), " ==> ", "key.txt")
		_ = ioutil.WriteFile("key.txt", mkey, 0644)
		os.Exit(0)
	case "upload":
		loadKey()
		initKey()

		if path != "" {
			url := ""
			if qpath == "" {
				url = UploadDir(path)
			} else {
				url = UploadDir(path, qpath)
			}
			fmt.Println("url:", url)
		} else {
			flag.Usage()
		}
	case "list":
		loadKey()
		initKey()
		for {
			data, err := getTencentYunCosFiles(path)
			if err != nil {
				fmt.Println("错误 106:", err)
				return
			}
			//for e := range data.CommonPrefixes {
			//	fmt.Println(data.CommonPrefixes[e])
			//}
			var files []string
			for e := range data.Contents {
				files = append(files, data.Contents[e].Key)
			}
			_, menu := askSelect("Path:"+path, append([]string{".."}, append(data.CommonPrefixes, files...)...)...)
			if menu == ".." {
				pathSplit := strings.Split(path, "/")
				if len(pathSplit) > 1 {
					path = strings.Join(pathSplit[0:len(pathSplit)-1], "/")
				} else {
					path = ""
				}
			} else {
				path = menu
			}
		}

	default:
		flag.Usage()
	}
}

type TencentYunConfig struct {
	SecretID     string
	SecretKey    string
	SessionToken string
	BucketURL    string
	Root         string
}

var client *cos.Client

func getTencentYunConf(c *TencentYunConfig) *cos.Client {
	if client == nil {
		BucketURL, _ := url.Parse(c.BucketURL)
		b := &cos.BaseURL{
			BucketURL: BucketURL,
		}
		// 1.永久密钥
		client = cos.NewClient(b, &http.Client{
			Transport: &cos.AuthorizationTransport{
				SecretID:  c.SecretID,
				SecretKey: c.SecretKey,
			},
		})
	}

	return client
}

func getTencentYunCosFiles(prefix string) (*cos.BucketGetResult, error) {

	// 1.永久密钥
	client := getTencentYunConf(&config)
	if prefix == "" {
		prefix = strings.Replace(config.Root+"/", `//`, `/`, -1)
	}
	res, _, err := client.Bucket.Get(context.Background(), &cos.BucketGetOptions{
		Prefix:    prefix,
		Delimiter: "/",
	})
	if err != nil {
		return nil, err
	}
	return res, nil
}

func UploadDir(fpath string, qpaths ...string) string {
	qpath := ""
	if len(qpaths) > 0 {
		qpath = qpaths[0]
	} else {
		separator := "/"
		if runtime.GOOS == "windows" {
			separator = `\`
		}
		separatordir := strings.Split(fpath, separator)
		if len(separatordir) > 0 {
			qpath = separatordir[len(separatordir)-1]
		}
	}
	client := getTencentYunConf(&config)

	files, _ := ioutil.ReadDir(fpath)
	for _, f := range files {
		if f.Name()[0] == '.' {
			continue
		}
		if f.IsDir() {
			UploadDir(filepath.Join(fpath, f.Name()), filepath.Join(qpath, f.Name()))
		} else {
			filename := filepath.Join(fpath, f.Name())
			fmt.Print("上传:", filename, " ==> ", filepath.Join(config.Root, filename), "   ")
			_, err := client.Object.PutFromFile(context.Background(), strings.Replace(filepath.Join(config.Root, qpath, f.Name()), `\`, `/`, -1), filename, nil)
			if err != nil {
				fmt.Println("错误 105:", err)
			} else {
				fmt.Println("完成")
			}
		}
	}

	urlPath := "/" + strings.Replace(filepath.Join(config.Root, qpath), `\`, `/`, -1)
	for strings.Index(urlPath, "//") > -1 {
		urlPath = strings.Replace(urlPath, "//", "/", -1)
	}

	return config.BucketURL + urlPath
}

func askSelect(title string, v ...string) (int, string) {

	result := -1
	prompt := &survey.Select{
		Message: title,
		Options: v,
	}

	err := survey.AskOne(prompt, &result)
	if err != nil {
		return 0, ""
	}

	if result == -1 {
		return askSelect(title, v...)
	}
	return result, v[result]
}

func askText(title string) string {
	result := ""
	prompt := &survey.Input{
		Message: title,
	}
	survey.AskOne(prompt, &result)

	if result == "" {
		fmt.Println("不允许为空")
		return askText(title)
	}
	return result
}
func home() (string, error) {
	user, err := user.Current()
	if nil == err {
		return user.HomeDir, nil
	}

	// cross compile support

	if "windows" == runtime.GOOS {
		return homeWindows()
	}

	// Unix-like system, so just assume Unix
	return homeUnix()
}

func homeUnix() (string, error) {
	// First prefer the HOME environmental variable
	if home := os.Getenv("HOME"); home != "" {
		return home, nil
	}

	// If that fails, try the shell
	var stdout bytes.Buffer
	cmd := exec.Command("sh", "-c", "eval echo ~$USER")
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		return "", err
	}

	result := strings.TrimSpace(stdout.String())
	if result == "" {
		return "", errors.New("blank output when reading home directory")
	}

	return result, nil
}

func homeWindows() (string, error) {
	drive := os.Getenv("HOMEDRIVE")
	path := os.Getenv("HOMEPATH")
	home := drive + path
	if drive == "" || path == "" {
		home = os.Getenv("USERPROFILE")
	}
	if home == "" {
		return "", errors.New("HOMEDRIVE, HOMEPATH, and USERPROFILE are blank")
	}

	return home, nil
}