From 35600e920002d904210f7bc37e9b7548263a2540 Mon Sep 17 00:00:00 2001 From: Tom Neuber Date: Thu, 25 Jan 2024 12:52:09 +0100 Subject: [PATCH] simplify downloader --- go.mod | 6 +- go.sum | 25 ++++- pkg/downloader/downloader.go | 186 ++++++----------------------------- 3 files changed, 56 insertions(+), 161 deletions(-) diff --git a/go.mod b/go.mod index c150c1b..eb7e133 100644 --- a/go.mod +++ b/go.mod @@ -4,12 +4,12 @@ go 1.21.6 require ( github.com/alecthomas/kong v0.8.1 - github.com/dustin/go-humanize v1.0.1 github.com/go-chi/chi/v5 v5.0.11 github.com/go-chi/render v1.0.3 github.com/hashicorp/go-memdb v1.3.4 github.com/levigross/grequests v0.0.0-20231203190023-9c307ef1f48d github.com/praserx/ipconv v1.2.1 + github.com/schollz/progressbar/v3 v3.14.1 ) require ( @@ -17,5 +17,9 @@ require ( github.com/google/go-querystring v1.1.0 // indirect github.com/hashicorp/go-immutable-radix v1.3.1 // indirect github.com/hashicorp/golang-lru v1.0.2 // indirect + github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect + github.com/rivo/uniseg v0.4.4 // indirect golang.org/x/net v0.19.0 // indirect + golang.org/x/sys v0.16.0 // indirect + golang.org/x/term v0.16.0 // indirect ) diff --git a/go.sum b/go.sum index 0be8204..36971b9 100644 --- a/go.sum +++ b/go.sum @@ -6,8 +6,9 @@ github.com/alecthomas/kong v0.8.1 h1:acZdn3m4lLRobeh3Zi2S2EpnXTd1mOL6U7xVml+vfkY github.com/alecthomas/kong v0.8.1/go.mod h1:n1iCIO2xS46oE8ZfYCNDqdR0b0wZNrXAIAqro/2132U= github.com/alecthomas/repr v0.1.0 h1:ENn2e1+J3k09gyj2shc0dHr/yjaWSHRlrJ4DPMevDqE= github.com/alecthomas/repr v0.1.0/go.mod h1:2kn6fqh/zIyPLmm3ugklbEi5hg5wS435eygvNfaDQL8= -github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= -github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/go-chi/chi/v5 v5.0.11 h1:BnpYbFZ3T3S1WMpD79r7R5ThWX40TaFB7L31Y8xqSwA= github.com/go-chi/chi/v5 v5.0.11/go.mod h1:DslCQbL2OYiznFReuXYUmQ2hGd1aDpCnlMNITLSKoi8= github.com/go-chi/render v1.0.3 h1:AsXqd2a1/INaIfUSKq3G5uA8weYx20FOsM7uSoCyyt4= @@ -29,10 +30,30 @@ github.com/hashicorp/golang-lru v1.0.2 h1:dV3g9Z/unq5DpblPpw+Oqcv4dU/1omnb4Ok8iP github.com/hashicorp/golang-lru v1.0.2/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4= github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM= github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg= +github.com/k0kubun/go-ansi v0.0.0-20180517002512-3bf9e2903213/go.mod h1:vNUNkEQ1e29fT/6vq2aBdFsgNPmy8qMdSay1npru+Sw= github.com/levigross/grequests v0.0.0-20231203190023-9c307ef1f48d h1:8fVmm2qScPn4JAF/YdTtqrPP3n58FgZ4GbKTNfaPuRs= github.com/levigross/grequests v0.0.0-20231203190023-9c307ef1f48d/go.mod h1:dFu6nuJHC3u9kCDcyGrEL7LwhK2m6Mt+alyiiIjDrRY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db h1:62I3jR2EmQ4l5rM/4FEfDWcRD+abF5XlKShorW5LRoQ= +github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db/go.mod h1:l0dey0ia/Uv7NcFFVbCLtqEBQbrT4OCwCSKTEv6enCw= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/praserx/ipconv v1.2.1 h1:MWGfrF+OZ0pqIuTlNlMgvJDDbohC3h751oN1+Ov3x4k= github.com/praserx/ipconv v1.2.1/go.mod h1:DSy+AKre/e3w/npsmUDMio+OR/a2rvmMdI7rerOIgqI= +github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis= +github.com/rivo/uniseg v0.4.4/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/schollz/progressbar/v3 v3.14.1 h1:VD+MJPCr4s3wdhTc7OEJ/Z3dAeBzJ7yKH/P4lC5yRTI= +github.com/schollz/progressbar/v3 v3.14.1/go.mod h1:Zc9xXneTzWXF81TGoqL71u0sBPjULtEHYtj/WVgVy8E= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= golang.org/x/net v0.19.0 h1:zTwKpTd2XuCqf8huc7Fo2iSy+4RHPd10s4KzeTnVr1c= golang.org/x/net v0.19.0/go.mod h1:CfAk/cbD4CthTvqiEl8NpboMuiuOYsAr/7NOjZJtv1U= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.14.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU= +golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/term v0.14.0/go.mod h1:TySc+nGkYR6qt8km8wUhuFRTVSMIX3XPR58y2lC8vww= +golang.org/x/term v0.16.0 h1:m+B6fahuftsE9qjo0VWp2FW0mB3MTJvR0BaMQrq0pmE= +golang.org/x/term v0.16.0/go.mod h1:yn7UURbUtPyrVJPGPq404EukNFxcm/foM+bV/bfcDsY= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= diff --git a/pkg/downloader/downloader.go b/pkg/downloader/downloader.go index fb5a06f..62f2127 100644 --- a/pkg/downloader/downloader.go +++ b/pkg/downloader/downloader.go @@ -1,198 +1,68 @@ package downloader import ( - "crypto/md5" - "crypto/sha1" - "crypto/sha256" "errors" "fmt" - "hash" "io" "os" "strings" - "time" - "github.com/dustin/go-humanize" req "github.com/levigross/grequests" + "github.com/schollz/progressbar/v3" ) -var ( - hashes = map[string]hash.Hash{ - "md5": md5.New(), //nolint: gosec - "sha1": sha1.New(), //nolint: gosec - "sha256": sha256.New(), - } -) - -type Context struct { - Filename string - Filesize int64 - Reader io.Reader - Chan chan *Response - Closer []func() error - hashWriter hash.Hash -} - -func NewContext(filename string, filesize int64, hw hash.Hash, closer ...func() error) *Context { - return &Context{ - Filename: filename, - Filesize: filesize, - Chan: make(chan *Response), - Closer: closer, - hashWriter: hw, - } -} - -func (ctx *Context) Checksum() string { - return fmt.Sprintf("%x", ctx.hashWriter.Sum(nil)) -} - -func (ctx *Context) SetReader(rdr io.Reader) { - ctx.Reader = io.TeeReader( - rdr, - io.MultiWriter( - ctx.hashWriter, - &ProgressWriter{ - Filename: ctx.Filename, - Filesize: ctx.Filesize, - Channel: ctx.Chan, - Start: time.Now(), - }, - ), - ) -} - -func (ctx *Context) Close() { - for i := len(ctx.Closer) - 1; i >= 0; i-- { - ctx.Closer[i]() - } -} - -type ProgressWriter struct { - Channel chan *Response - Start time.Time - Filename string - Filesize int64 - Bytes int64 -} - -func (pw *ProgressWriter) Write(data []byte) (int, error) { - pw.Bytes = pw.Bytes + int64(len(data)) - var progress float64 - if pw.Filesize > 0 { - progress = float64(pw.Bytes) / float64(pw.Filesize) * 100 - } - pw.Channel <- &Response{ - Progress: progress, - BytesPerSecond: float64(pw.Bytes) / time.Since(pw.Start).Seconds(), - BytesTransfered: pw.Bytes, - BytesTotal: pw.Filesize, - } - return len(data), nil -} - -type File struct { - URL string - Filename string - Size int64 -} - -type Response struct { - Error error - Progress float64 - BytesPerSecond float64 - BytesTransfered int64 - BytesTotal int64 -} - -func handleError(err error, adds ...string) { - msg := err.Error() - if len(adds) != 0 { - msg = fmt.Sprintf("%s %s", msg, strings.Join(adds, " ")) - } - fmt.Println(msg) +func handleError(err error) { + fmt.Println(err.Error()) os.Exit(0) } -func printProgress(ctx *Context) { - for { - resp, ok := <-ctx.Chan - if !ok { - fmt.Println("\nSaved file to", ctx.Filename) - for key, value := range hashes { - fmt.Printf("%s: %x\n", key, value.Sum(nil)) - } - break - } - if resp.Error != nil { - handleError(resp.Error) - } - if resp.BytesTotal > 0 { - fmt.Printf("\rDownloading %7s/%7s %.02f%%", humanize.Bytes(uint64(resp.BytesTransfered)), humanize.Bytes(uint64(resp.BytesTotal)), resp.Progress) - } else { - fmt.Printf("\rDownloading %7s", humanize.Bytes(uint64(resp.BytesTransfered))) - } - } -} - func DownloadFile(path, url string) { - file := File{ - URL: url, - Filename: path, - } - response, err := req.Head(file.URL, nil) + resp, err := req.Head(url, nil) if err != nil { if strings.Contains(err.Error(), "no such host") { handleError(errors.New("invalid url")) } else if strings.Contains(err.Error(), "certificate signed by unknown authority") { - handleError(errors.New("certificate signed by unknown authority")) + handleError(errors.New("certificate from unknown authority")) } handleError(err) } - defer response.Close() - if response.StatusCode == 404 { + defer resp.Close() + if resp.StatusCode == 404 { handleError(errors.New("invalid url")) + } else if resp.StatusCode == 401 { + handleError(errors.New("restricted access (credentials required)")) + } else if resp.StatusCode != 200 { + handleError(errors.New(resp.RawResponse.Status)) } - if response.StatusCode == 401 { - handleError(errors.New("restriced access, credentials required")) - } - if response.StatusCode != 200 { - handleError(errors.New(response.RawResponse.Status)) - } - response.Close() + resp.Close() - resp, err := req.Get(file.URL, nil) + resp, err = req.Get(url, nil) if err != nil { handleError(err) } var filesize int64 - if resp.RawResponse.ContentLength != -1 { + if resp.RawResponse.ContentLength > -1 { filesize = resp.RawResponse.ContentLength } - ctx := NewContext(file.Filename, filesize, md5.New()) - ctx.SetReader(resp.RawResponse.Body) - go func(ctx *Context) { - defer ctx.Close() + destFile, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0644) + if err != nil { + handleError(err) + } + defer destFile.Close() - dstFile, err := os.Create(ctx.Filename) - if err != nil { - ctx.Chan <- &Response{Error: err} - close(ctx.Chan) - return - } - defer dstFile.Close() + bar := progressbar.DefaultBytes( + filesize, + "downloading", + ) - writer := io.MultiWriter(dstFile, hashes["md5"], hashes["sha1"], hashes["sha256"]) - _, err = io.Copy(writer, ctx.Reader) - if err != nil { - ctx.Chan <- &Response{Error: err} - } - close(ctx.Chan) - }(ctx) + _, err = io.Copy(io.MultiWriter(destFile, bar), resp.RawResponse.Body) + if err != nil { + handleError(err) + } - printProgress(ctx) + fmt.Println("Saved file to", path) } func FileExists(path string) bool {