ufile/cmd/ufile-cli/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

var (
	version = "dev"
	commit  = "unknown"
	date    = "unknown"
)

const (
	exitOK          = 0
	exitGeneral     = 1
	exitUsage       = 2
	exitAuth        = 3
	exitPermission  = 4
	exitUnavailable = 5
	exitValidation  = 6
	exitNetwork     = 7
	exitLocalIO     = 8
)

type streams struct {
	in     io.Reader
	out    io.Writer
	errOut io.Writer
	stdin  *os.File
}

type globalOptions struct {
	host       string
	configPath string
	timeout    time.Duration
}

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	os.Exit(run(ctx, os.Args[1:], streams{in: os.Stdin, out: os.Stdout, errOut: os.Stderr, stdin: os.Stdin}))
}

func run(ctx context.Context, args []string, ioStreams streams) int {
	options, rest, help, err := parseGlobalOptions(args)
	if err != nil {
		return writeCommandError(ioStreams.errOut, usageError(err.Error()))
	}
	if help || len(rest) == 0 {
		_, _ = io.WriteString(ioStreams.out, rootHelp)
		return exitOK
	}
	if rest[0] == "help" {
		topic := ""
		if len(rest) > 1 {
			topic = strings.Join(rest[1:], " ")
		}
		return printHelp(ioStreams.out, ioStreams.errOut, topic)
	}
	if rest[0] == "version" {
		return writeJSON(ioStreams.out, map[string]string{"version": version, "commit": commit, "date": date})
	}
	if helpTopic, ok := requestedCommandHelp(rest); ok {
		return printHelp(ioStreams.out, ioStreams.errOut, helpTopic)
	}
	if err := validateCommand(rest); err != nil {
		return writeCommandError(ioStreams.errOut, err)
	}

	config, err := loadConfig(options.configPath)
	if err != nil {
		return writeCommandError(ioStreams.errOut, localError("读取配置失败", err))
	}
	host := firstNonEmpty(options.host, os.Getenv("UFILE_HOST"), os.Getenv("UFILE_SERVER"), config.Host)
	if host == "" {
		return writeCommandError(ioStreams.errOut, usageError("缺少服务地址；请使用 -H/--host，或设置 UFILE_HOST"))
	}
	host, err = normalizeHost(host)
	if err != nil {
		return writeCommandError(ioStreams.errOut, usageError(err.Error()))
	}
	session := firstNonEmpty(os.Getenv("UFILE_SESSION"), config.Session)
	client := newAPIClient(host, session, options.timeout)

	var commandErr *commandError
	switch rest[0] {
	case "auth":
		commandErr = runAuth(ctx, rest[1:], options.configPath, config, client, ioStreams)
	case "upload":
		commandErr = runUpload(ctx, rest[1:], client, ioStreams)
	case "download":
		commandErr = runDownload(ctx, rest[1:], client, ioStreams)
	default:
		commandErr = usageError(fmt.Sprintf("未知命令 %q；运行 ufile-cli help 查看可用命令", rest[0]))
	}
	if commandErr != nil {
		return writeCommandError(ioStreams.errOut, commandErr)
	}
	return exitOK
}

func parseGlobalOptions(args []string) (globalOptions, []string, bool, error) {
	var options globalOptions
	var help, showVersion bool
	fs := flag.NewFlagSet("ufile-cli", flag.ContinueOnError)
	fs.SetOutput(io.Discard)
	fs.StringVar(&options.host, "H", "", "ufile 服务地址")
	fs.StringVar(&options.host, "host", "", "ufile 服务地址")
	fs.StringVar(&options.host, "server", "", "ufile 服务地址（--host 的别名）")
	fs.StringVar(&options.configPath, "config", defaultConfigPath(), "配置文件路径")
	fs.DurationVar(&options.timeout, "timeout", 10*time.Minute, "单次请求超时")
	fs.BoolVar(&help, "help", false, "显示帮助")
	fs.BoolVar(&showVersion, "V", false, "显示版本")
	fs.BoolVar(&showVersion, "version", false, "显示版本")
	if err := fs.Parse(args); err != nil {
		return options, nil, false, err
	}
	if options.timeout < 0 {
		return options, nil, false, fmt.Errorf("--timeout 不能为负数")
	}
	if showVersion {
		return options, []string{"version"}, false, nil
	}
	return options, fs.Args(), help, nil
}

func firstNonEmpty(values ...string) string {
	for _, value := range values {
		if value != "" {
			return value
		}
	}
	return ""
}

func interspersedArgs(args []string, valueFlags map[string]bool) []string {
	options := make([]string, 0, len(args))
	positionals := make([]string, 0, len(args))
	for index := 0; index < len(args); index++ {
		argument := args[index]
		if argument == "--" {
			positionals = append(positionals, args[index+1:]...)
			break
		}
		if argument == "-" || !strings.HasPrefix(argument, "-") {
			positionals = append(positionals, argument)
			continue
		}
		options = append(options, argument)
		name := strings.TrimLeft(strings.SplitN(argument, "=", 2)[0], "-")
		if valueFlags[name] && !strings.Contains(argument, "=") && index+1 < len(args) {
			index++
			options = append(options, args[index])
		}
	}
	return append(options, positionals...)
}

func requestedCommandHelp(args []string) (string, bool) {
	if args[0] == "auth" && len(args) > 1 && args[1] == "help" {
		return "auth", true
	}
	for _, argument := range args[1:] {
		if argument == "--help" {
			topic := args[0]
			if topic == "auth" && len(args) > 1 && !strings.HasPrefix(args[1], "-") {
				topic += " " + args[1]
			}
			return topic, true
		}
	}
	return "", false
}

func validateCommand(args []string) *commandError {
	switch args[0] {
	case "upload", "download":
		return nil
	case "auth":
		if len(args) < 2 {
			return usageError("auth 需要子命令 login、whoami 或 logout；运行 ufile-cli help auth")
		}
		switch args[1] {
		case "login", "whoami", "logout", "help":
			return nil
		default:
			return usageError(fmt.Sprintf("未知 auth 子命令 %q；运行 ufile-cli help auth", args[1]))
		}
	}
	return usageError(fmt.Sprintf("未知命令 %q；运行 ufile-cli help 查看可用命令", args[0]))
}