237 lines
7.1 KiB
Go
237 lines
7.1 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"math/big"
|
|
"mime"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"net/textproto"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"regexp"
|
|
"strconv"
|
|
"time"
|
|
)
|
|
|
|
var pickupCodePattern = regexp.MustCompile(`^[0-9]{6}$`)
|
|
|
|
type uploadResult struct {
|
|
Hash string `json:"hash"`
|
|
Name string `json:"name"`
|
|
Size int64 `json:"size"`
|
|
ExpiresAt *time.Time `json:"expiresAt"`
|
|
MaxDownloads *int `json:"maxDownloads"`
|
|
PickupCode string `json:"pickupCode,omitempty"`
|
|
DownloadURL string `json:"downloadUrl"`
|
|
PreviewURL string `json:"previewUrl"`
|
|
}
|
|
|
|
func runUpload(ctx context.Context, args []string, client *apiClient, ioStreams streams) *commandError {
|
|
var retention, downloads, pickupCode, name, contentType string
|
|
var noPickupCode, help bool
|
|
fs := flag.NewFlagSet("upload", flag.ContinueOnError)
|
|
fs.SetOutput(io.Discard)
|
|
fs.StringVar(&retention, "R", "1d", "保留时间")
|
|
fs.StringVar(&retention, "retention", "1d", "保留时间")
|
|
fs.StringVar(&downloads, "D", "5", "下载次数")
|
|
fs.StringVar(&downloads, "downloads", "5", "下载次数")
|
|
fs.StringVar(&pickupCode, "P", "", "6 位取件码")
|
|
fs.StringVar(&pickupCode, "pickup-code", "", "6 位取件码")
|
|
fs.BoolVar(&noPickupCode, "no-pickup-code", false, "不设置取件码")
|
|
fs.StringVar(&name, "N", "", "标准输入的文件名")
|
|
fs.StringVar(&name, "name", "", "标准输入的文件名")
|
|
fs.StringVar(&contentType, "content-type", "", "MIME 类型")
|
|
fs.BoolVar(&help, "help", false, "显示帮助")
|
|
args = interspersedArgs(args, map[string]bool{"R": true, "retention": true, "D": true, "downloads": true, "P": true, "pickup-code": true, "N": true, "name": true, "content-type": true})
|
|
if err := fs.Parse(args); err != nil {
|
|
return usageError(err.Error())
|
|
}
|
|
if help {
|
|
_, _ = io.WriteString(ioStreams.out, uploadHelp)
|
|
return nil
|
|
}
|
|
if fs.NArg() != 1 {
|
|
return usageError("upload 需要且只接受一个文件路径；运行 ufile-cli help upload")
|
|
}
|
|
if client.session == "" {
|
|
return &commandError{Code: "not_logged_in", Message: "上传需要登录；请先运行 auth login", ExitCode: exitAuth}
|
|
}
|
|
if !validRetention(retention) {
|
|
return usageError("--retention 只支持 1h、8h、1d、3d 或 forever")
|
|
}
|
|
if !validDownloads(downloads) {
|
|
return usageError("--downloads 需要是 1-99 或 unlimited")
|
|
}
|
|
if pickupCode != "" && noPickupCode {
|
|
return usageError("--pickup-code 与 --no-pickup-code 不能同时使用")
|
|
}
|
|
if pickupCode != "" && !pickupCodePattern.MatchString(pickupCode) {
|
|
return usageError("--pickup-code 必须是 6 位数字")
|
|
}
|
|
if pickupCode == "" && !noPickupCode {
|
|
var err error
|
|
pickupCode, err = randomPickupCode()
|
|
if err != nil {
|
|
return &commandError{Code: "random_failed", Message: "生成取件码失败", ExitCode: exitGeneral, Cause: err}
|
|
}
|
|
}
|
|
|
|
path := fs.Arg(0)
|
|
var source io.Reader
|
|
var closeSource func() error
|
|
if path == "-" {
|
|
if name == "" {
|
|
return usageError("从标准输入上传时必须使用 -N/--name 指定文件名")
|
|
}
|
|
source = ioStreams.in
|
|
} else {
|
|
if name != "" {
|
|
return usageError("-N/--name 只用于标准输入上传")
|
|
}
|
|
file, err := os.Open(path)
|
|
if err != nil {
|
|
return localError("打开待上传文件失败", err)
|
|
}
|
|
info, err := file.Stat()
|
|
if err != nil {
|
|
file.Close()
|
|
return localError("读取待上传文件信息失败", err)
|
|
}
|
|
if !info.Mode().IsRegular() {
|
|
file.Close()
|
|
return usageError("upload 只支持普通文件或标准输入")
|
|
}
|
|
source = file
|
|
closeSource = file.Close
|
|
name = filepath.Base(path)
|
|
}
|
|
if closeSource != nil {
|
|
defer closeSource()
|
|
}
|
|
name = filepath.Base(name)
|
|
if name == "" || name == "." {
|
|
return usageError("文件名不能为空")
|
|
}
|
|
if contentType == "" {
|
|
contentType = mime.TypeByExtension(filepath.Ext(name))
|
|
}
|
|
if contentType == "" {
|
|
contentType = "application/octet-stream"
|
|
}
|
|
if _, _, err := mime.ParseMediaType(contentType); err != nil {
|
|
return usageError("--content-type 不是有效的 MIME 类型")
|
|
}
|
|
|
|
bodyReader, bodyWriter := io.Pipe()
|
|
multipartWriter := multipart.NewWriter(bodyWriter)
|
|
writeResult := make(chan error, 1)
|
|
go func() {
|
|
err := writeMultipartUpload(multipartWriter, source, name, contentType, retention, downloads, pickupCode)
|
|
if err == nil {
|
|
err = multipartWriter.Close()
|
|
}
|
|
_ = bodyWriter.CloseWithError(err)
|
|
writeResult <- err
|
|
}()
|
|
|
|
request, err := client.newRequest(ctx, http.MethodPost, "/api/files", bodyReader)
|
|
if err != nil {
|
|
bodyReader.CloseWithError(err)
|
|
<-writeResult
|
|
return classifyError(err)
|
|
}
|
|
request.Header.Set("Content-Type", multipartWriter.FormDataContentType())
|
|
response, requestErr := client.http.Do(request)
|
|
if requestErr != nil {
|
|
bodyReader.CloseWithError(requestErr)
|
|
<-writeResult
|
|
return classifyError(requestErr)
|
|
}
|
|
_ = bodyReader.Close()
|
|
writeErr := <-writeResult
|
|
defer response.Body.Close()
|
|
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
|
return classifyError(decodeAPIError(response))
|
|
}
|
|
if writeErr != nil {
|
|
return localError("读取待上传文件失败", writeErr)
|
|
}
|
|
var uploaded fileResponse
|
|
if err := json.NewDecoder(io.LimitReader(response.Body, 2<<20)).Decode(&uploaded); err != nil {
|
|
return networkError("服务端返回了无效的上传结果", err)
|
|
}
|
|
downloadURL := shareURL(client.host, uploaded.Hash, pickupCode, false)
|
|
result := uploadResult{
|
|
Hash: uploaded.Hash, Name: uploaded.Name, Size: uploaded.Size, ExpiresAt: uploaded.ExpiresAt,
|
|
MaxDownloads: uploaded.MaxDownloads, PickupCode: pickupCode, DownloadURL: downloadURL,
|
|
PreviewURL: shareURL(client.host, uploaded.Hash, pickupCode, true),
|
|
}
|
|
writeJSON(ioStreams.out, result)
|
|
return nil
|
|
}
|
|
|
|
func writeMultipartUpload(writer *multipart.Writer, source io.Reader, name, contentType, retention, downloads, pickupCode string) error {
|
|
fields := [][2]string{{"retention", retention}, {"downloads", downloads}, {"pickupCode", pickupCode}}
|
|
for _, field := range fields {
|
|
if err := writer.WriteField(field[0], field[1]); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
header := make(textproto.MIMEHeader)
|
|
header.Set("Content-Disposition", mime.FormatMediaType("form-data", map[string]string{"name": "file", "filename": name}))
|
|
header.Set("Content-Type", contentType)
|
|
part, err := writer.CreatePart(header)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, err = io.Copy(part, source)
|
|
return err
|
|
}
|
|
|
|
func randomPickupCode() (string, error) {
|
|
value, err := rand.Int(rand.Reader, big.NewInt(1_000_000))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return fmt.Sprintf("%06d", value.Int64()), nil
|
|
}
|
|
|
|
func validRetention(value string) bool {
|
|
switch value {
|
|
case "1h", "8h", "1d", "3d", "forever":
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func validDownloads(value string) bool {
|
|
if value == "unlimited" {
|
|
return true
|
|
}
|
|
number, err := strconv.Atoi(value)
|
|
return err == nil && number >= 1 && number <= 99
|
|
}
|
|
|
|
func shareURL(host, hash, pickupCode string, preview bool) string {
|
|
result := host + "/uf/" + url.PathEscape(hash)
|
|
query := url.Values{}
|
|
if pickupCode != "" {
|
|
query.Set("code", pickupCode)
|
|
}
|
|
if preview {
|
|
query.Set("preview", "true")
|
|
}
|
|
if encoded := query.Encode(); encoded != "" {
|
|
result += "?" + encoded
|
|
}
|
|
return result
|
|
}
|