ufile/cmd/ufile-cli/download.go

package main

import (
	"context"
	"errors"
	"flag"
	"io"
	"mime"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
)

type downloadResult struct {
	Path        string `json:"path"`
	Name        string `json:"name"`
	Size        int64  `json:"size"`
	ContentType string `json:"contentType"`
}

func runDownload(ctx context.Context, args []string, client *apiClient, ioStreams streams) *commandError {
	var code, output string
	var force, help bool
	fs := flag.NewFlagSet("download", flag.ContinueOnError)
	fs.SetOutput(io.Discard)
	fs.StringVar(&code, "C", "", "取件码")
	fs.StringVar(&code, "code", "", "取件码")
	fs.StringVar(&output, "O", "", "输出路径")
	fs.StringVar(&output, "output", "", "输出路径")
	fs.BoolVar(&force, "force", false, "覆盖已有文件")
	fs.BoolVar(&help, "help", false, "显示帮助")
	args = interspersedArgs(args, map[string]bool{"C": true, "code": true, "O": true, "output": true})
	if err := fs.Parse(args); err != nil {
		return usageError(err.Error())
	}
	if help {
		_, _ = io.WriteString(ioStreams.out, downloadHelp)
		return nil
	}
	if fs.NArg() != 1 {
		return usageError("download 需要且只接受一个 8 位文件码；运行 ufile-cli help download")
	}
	hash := strings.TrimSpace(fs.Arg(0))
	if len(hash) != 8 {
		return usageError("文件码必须是 8 位")
	}
	if code != "" && !pickupCodePattern.MatchString(code) {
		return usageError("--code 必须是 6 位数字")
	}
	if output == "" {
		return usageError("必须使用 -O/--output 指定输出路径；使用 - 可写入 stdout")
	}
	if output == "-" && force {
		return usageError("输出到 stdout 时不能使用 --force")
	}

	var destination io.Writer = ioStreams.out
	var temp *os.File
	var tempName string
	if output != "-" {
		if info, err := os.Stat(output); err == nil && info.IsDir() {
			return localError("目标路径是目录，请使用完整文件路径", nil)
		} else if err == nil && !force {
			return localError("目标文件已存在；如需覆盖请使用 --force", nil)
		} else if err != nil && !errors.Is(err, os.ErrNotExist) {
			return localError("检查目标文件失败", err)
		}
		var err error
		temp, err = os.CreateTemp(filepath.Dir(output), ".ufile-download-*")
		if err != nil {
			return localError("创建临时下载文件失败", err)
		}
		tempName = temp.Name()
		defer os.Remove(tempName)
		destination = temp
	}

	query := url.Values{}
	if code != "" {
		query.Set("code", code)
	}
	path := "/uf/" + url.PathEscape(hash)
	if encoded := query.Encode(); encoded != "" {
		path += "?" + encoded
	}
	request, err := client.newRequest(ctx, http.MethodGet, path, nil)
	if err != nil {
		closeTemp(temp)
		return classifyError(err)
	}
	request.Header.Set("Accept", "application/octet-stream")
	response, err := client.http.Do(request)
	if err != nil {
		closeTemp(temp)
		return classifyError(err)
	}
	defer response.Body.Close()
	if response.StatusCode < 200 || response.StatusCode >= 300 {
		closeTemp(temp)
		return classifyError(decodeAPIError(response))
	}

	size, copyErr := io.Copy(destination, response.Body)
	if temp != nil {
		closeErr := temp.Close()
		temp = nil
		if copyErr == nil {
			copyErr = closeErr
		}
	}
	if copyErr != nil {
		return localError("写入下载内容失败；本次访问次数可能已经被消耗", copyErr)
	}
	if output == "-" {
		return nil
	}
	if force {
		if err := os.Remove(output); err != nil && !errors.Is(err, os.ErrNotExist) {
			return localError("覆盖已有文件失败；本次访问次数已经被消耗", err)
		}
	} else if _, err := os.Stat(output); err == nil {
		return localError("下载期间目标文件已被创建，未覆盖该文件；本次访问次数已经被消耗", nil)
	} else if err != nil && !errors.Is(err, os.ErrNotExist) {
		return localError("检查目标文件失败；本次访问次数已经被消耗", err)
	}
	if err := os.Rename(tempName, output); err != nil {
		return localError("保存下载文件失败；本次访问次数已经被消耗", err)
	}

	result := downloadResult{
		Path: output, Name: responseFilename(response), Size: size,
		ContentType: response.Header.Get("Content-Type"),
	}
	writeJSON(ioStreams.out, result)
	return nil
}

func responseFilename(response *http.Response) string {
	_, params, err := mime.ParseMediaType(response.Header.Get("Content-Disposition"))
	if err != nil {
		return ""
	}
	return filepath.Base(params["filename"])
}

func closeTemp(file *os.File) {
	if file != nil {
		_ = file.Close()
	}
}