From 3403deb463f926a838431cbd8066670ac9568ab0 Mon Sep 17 00:00:00 2001 From: Leon Mika Date: Sat, 4 Jan 2025 14:21:10 +1100 Subject: [PATCH] Fixed usage and progress display --- .idea/vcs.xml | 6 +++++ chunker.go | 20 +++++++++++----- config.go | 2 ++ display.go | 64 +++++++++++++++++++++++++++++++++++++++++++++++++++ main.go | 57 +++++++++++++++++++++++++++++---------------- 5 files changed, 123 insertions(+), 26 deletions(-) create mode 100644 .idea/vcs.xml create mode 100644 display.go diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 0000000..94a25f7 --- /dev/null +++ b/.idea/vcs.xml @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/chunker.go b/chunker.go index d72d8ba..1115391 100644 --- a/chunker.go +++ b/chunker.go @@ -17,13 +17,17 @@ type chunker struct { gc *gokapiClient parallelChunks int chunkSize int + maxDownloads int + maxDays int } -func newChunker(gc *gokapiClient, parallelChunks, chunkSize int) *chunker { +func newChunker(gc *gokapiClient, config Config) *chunker { return &chunker{ gc: gc, - parallelChunks: parallelChunks, - chunkSize: chunkSize, + parallelChunks: config.ParallelChunks, + chunkSize: config.ChunkSize, + maxDownloads: config.MaxDownloads, + maxDays: config.MaxDays, } } @@ -45,8 +49,8 @@ func (c *chunker) UploadFile(ctx context.Context, filename string, progress func filename: fname, totalSize: fstat.Size(), contentType: mime.TypeByExtension(filepath.Ext(fname)), - allowedDownloads: 5, - expiryDays: 7, + allowedDownloads: c.maxDownloads, + expiryDays: c.maxDays, password: "", } @@ -61,6 +65,10 @@ func (c *chunker) upload(ctx context.Context, fi uploadInfo, r io.ReaderAt, prog } chunks := int(fi.totalSize/int64(c.chunkSize) + 1) + uploaders := c.parallelChunks + if chunks < uploaders { + uploaders = 1 + } chunkUploaded := make(chan uploadedChunk) doneChunkReport := make(chan struct{}) @@ -91,7 +99,7 @@ func (c *chunker) upload(ctx context.Context, fi uploadInfo, r io.ReaderAt, prog }() errGroup, egctx := errgroup.WithContext(ctx) - errGroup.SetLimit(c.parallelChunks) + errGroup.SetLimit(uploaders) for i := 0; i < chunks; i++ { errGroup.Go(func() error { diff --git a/config.go b/config.go index a345d75..f77d5b5 100644 --- a/config.go +++ b/config.go @@ -5,4 +5,6 @@ type Config struct { APIKey string ParallelChunks int ChunkSize int + MaxDownloads int + MaxDays int } diff --git a/display.go b/display.go new file mode 100644 index 0000000..9e3ba01 --- /dev/null +++ b/display.go @@ -0,0 +1,64 @@ +package main + +import ( + "flag" + "fmt" + "github.com/schollz/progressbar/v3" + "os" + "path/filepath" + "strings" +) + +var ( + ansiMoveUp = []byte{keyEscape, '[', '1', 'A'} + ansiClearLine = []byte{keyEscape, '[', '2', 'K'} + ansiLineFeed = []byte{keyEscape, '[', '1', 'G'} +) + +func envVarUsage() { + fmt.Fprintf(flag.CommandLine.Output(), " %-20s Gokapi URL (e.g. https://gokapi.example.com/)\n", envVarURL) + fmt.Fprintf(flag.CommandLine.Output(), " %-20s Gokapi API key\n", envVarAPIKey) +} + +func init() { + flag.Usage = func() { + fmt.Fprintf(flag.CommandLine.Output(), "Usage: %s [FLAGS] FILES...\n", os.Args[0]) + flag.PrintDefaults() + fmt.Fprintf(flag.CommandLine.Output(), "Environment variables:\n") + envVarUsage() + } +} + +func progressBarReport(filename string, maxFileLen int) func(cr ChunkReport) { + var pr *progressbar.ProgressBar + + displayedFileName := formatFilename(filename, maxFileLen) + return func(cr ChunkReport) { + if cr.UploadedChunks == 0 && pr == nil { + pr = progressbar.DefaultBytes(cr.TotalSize, displayedFileName) + } + pr.Set(int(cr.UploadedBytes)) + } +} + +func displayCompletion(filename string, maxFileLen int, ur UploadResponse) { + os.Stdout.Write(ansiMoveUp) + os.Stdout.Write(ansiClearLine) + os.Stdout.Write(ansiLineFeed) + + displayedFileName := formatFilename(filename, maxFileLen) + fmt.Printf("%v: done - %v\n", displayedFileName, ur.FileInfo.UrlHotlink) +} + +func displayFailedCompletion(filename string, maxFileLen int, err error) { + os.Stdout.Write(ansiClearLine) + os.Stdout.Write(ansiLineFeed) + + displayedFileName := formatFilename(filename, maxFileLen) + fmt.Printf("%v: error - %v\n", displayedFileName, err) +} + +func formatFilename(filename string, maxFileLen int) string { + baseFileName := filepath.Base(filename) + return strings.Repeat(" ", maxFileLen-len(baseFileName)) + baseFileName +} diff --git a/main.go b/main.go index 3a147c9..46cc06c 100644 --- a/main.go +++ b/main.go @@ -3,21 +3,43 @@ package main import ( "context" "flag" - "github.com/schollz/progressbar/v3" + "fmt" "log" "net/url" "os" "path/filepath" ) +const ( + envVarURL = "SEND2GOKAPI_URL" + envVarAPIKey = "SEND2GOKAPI_API_KEY" + + keyEscape = 27 +) + func main() { + flagDownloads := flag.Int("d", 5, "Max number of downloads") + flagDays := flag.Int("t", 7, "Amount of time to keep file, in days") + flag.Parse() + + if flag.NArg() < 1 { + flag.Usage() + os.Exit(2) + } + config := Config{ - Hostname: os.Getenv("GOKAPI_HOSTNAME"), - APIKey: os.Getenv("GOKAPI_API_KEY"), + Hostname: os.Getenv(envVarURL), + APIKey: os.Getenv(envVarAPIKey), ParallelChunks: 4, ChunkSize: 1024 * 100, + MaxDownloads: max(*flagDownloads, 0), + MaxDays: max(*flagDays, 0), + } + if config.Hostname == "" || config.APIKey == "" { + fmt.Fprintf(flag.CommandLine.Output(), "Please set the following environment variables:\n") + envVarUsage() + os.Exit(2) } - flag.Parse() hostUrl, err := url.Parse(config.Hostname) if err != nil { @@ -27,24 +49,19 @@ func main() { ctx := context.Background() client := newGokapiClient(hostUrl, config.APIKey) - cnkr := newChunker(client, config.ParallelChunks, config.ChunkSize) + cnkr := newChunker(client, config) + + maxFileLen := 0 + for _, file := range flag.Args() { + maxFileLen = max(maxFileLen, len(filepath.Base(file))) + } for _, file := range flag.Args() { - _, err := cnkr.UploadFile(ctx, file, progressBarReport(file)) - if err != nil { - log.Fatal(err) + ur, err := cnkr.UploadFile(ctx, file, progressBarReport(file, maxFileLen)) + if err == nil { + displayCompletion(file, maxFileLen, ur) + } else { + displayFailedCompletion(file, maxFileLen, err) } } - -} - -func progressBarReport(filename string) func(cr ChunkReport) { - var pr *progressbar.ProgressBar - - return func(cr ChunkReport) { - if cr.UploadedChunks == 0 && pr == nil { - pr = progressbar.DefaultBytes(cr.TotalSize, filepath.Base(filename)) - } - pr.Set(int(cr.UploadedBytes)) - } }