113 lines
2.4 KiB
Go
113 lines
2.4 KiB
Go
package main
|
|
|
|
import (
|
|
"crypto/sha1"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
|
|
"github.com/go-resty/resty/v2"
|
|
"github.com/schollz/progressbar/v3"
|
|
)
|
|
|
|
type UploadCommand struct {
|
|
Url string `short:"u" long:"url" description:"Url of the instance" required:"true"`
|
|
Token string `short:"t" long:"token" description:"Token of the instance" required:"true"`
|
|
File string `short:"f" long:"file" description:"File to upload" required:"true"`
|
|
}
|
|
|
|
type UploadError struct {
|
|
Error string
|
|
}
|
|
|
|
type UploadChunkSuccess struct {
|
|
Checksum string
|
|
}
|
|
|
|
type UploadFileRequest struct {
|
|
Name string
|
|
Checksum string
|
|
Chunks []string
|
|
}
|
|
|
|
type UploadFileResponse struct {
|
|
Sum string
|
|
NbChunks uint32
|
|
Size uint64
|
|
}
|
|
|
|
func (c *UploadCommand) Execute(args []string) error {
|
|
f, err := os.Open(c.File)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer f.Close()
|
|
|
|
st, err := f.Stat()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
chunkSize := int64(1 << 20)
|
|
nbChunks := st.Size() / chunkSize
|
|
if st.Size()%chunkSize > 0 {
|
|
nbChunks++
|
|
}
|
|
|
|
uploadFile := &UploadFileRequest{Name: filepath.Base(c.File), Chunks: make([]string, nbChunks)}
|
|
progress := progressbar.DefaultBytes(st.Size(), fmt.Sprintf("Uploading %s", uploadFile.Name))
|
|
defer progress.Close()
|
|
|
|
cli := resty.New().SetBaseURL(c.Url).SetAuthScheme("Private").SetAuthToken(c.Token)
|
|
b := make([]byte, chunkSize)
|
|
checksum := sha1.New()
|
|
for i := 0; ; i++ {
|
|
n, err := f.Read(b)
|
|
if n == 0 {
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
return err
|
|
}
|
|
|
|
resp, err := cli.R().SetError(&UploadError{}).SetResult(&UploadChunkSuccess{}).SetBody(b[0:n]).Post("/file/chunk")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if err, ok := resp.Error().(*UploadError); ok {
|
|
return errors.New(err.Error)
|
|
}
|
|
|
|
if result, ok := resp.Result().(*UploadChunkSuccess); ok {
|
|
uploadFile.Chunks[i] = result.Checksum
|
|
checksum.Write(b[0:n])
|
|
progress.Add(n)
|
|
}
|
|
}
|
|
|
|
uploadFile.Checksum = hex.EncodeToString(checksum.Sum(nil))
|
|
|
|
resp, err := cli.R().SetBody(uploadFile).SetError(&UploadError{}).SetResult(&UploadFileResponse{}).Post("/file")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if err, ok := resp.Error().(*UploadError); ok {
|
|
logger.Println("Upload failed")
|
|
return errors.New(err.Error)
|
|
}
|
|
|
|
if result, ok := resp.Result().(*UploadFileResponse); ok {
|
|
fmt.Printf("Upload succeed\nSum: %s\nNbChunks: %d\nSize: %d\n", result.Sum, result.NbChunks, result.Size)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func init() {
|
|
parser.AddCommand("upload", "Upload a file", "", &UploadCommand{})
|
|
}
|