diff options
| -rw-r--r-- | fetcher.go | 68 |
1 files changed, 63 insertions, 5 deletions
@@ -27,11 +27,61 @@ package rhimport import ( "fmt" "github.com/golang-basic/go-curl" + "mime" "net/url" + "os" "path" + "strings" "time" ) +type CurlCBData struct { + basepath string + filename string + remotename string + *os.File +} + +func (self *CurlCBData) Cleanup() { + if self.File != nil { + self.File.Close() + } +} + +func curlHeaderCallback(ptr []byte, userdata interface{}) bool { + hdr := fmt.Sprintf("%s", ptr) + data := userdata.(*CurlCBData) + + if strings.HasPrefix(hdr, "Content-Disposition:") { + if mediatype, params, err := mime.ParseMediaType(strings.TrimPrefix(hdr, "Content-Disposition:")); err == nil { + if mediatype == "attachment" { + data.filename = data.basepath + "/" + params["filename"] + } + } + } + return true +} + +func curlWriteCallback(ptr []byte, userdata interface{}) bool { + data := userdata.(*CurlCBData) + if data.File == nil { + if data.filename == "" { + data.filename = data.basepath + "/" + data.remotename + } + fp, err := os.OpenFile(data.filename, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600) + if err != nil { + rhl.Printf("Unable to create file %s: %s", data.filename, err) + return false + } + data.File = fp + } + if _, err := data.File.Write(ptr); err != nil { + rhl.Printf("Unable to write file %s: %s", data.filename, err) + return false + } + return true +} + func FetchFileCurl(ctx *ImportContext, uri *url.URL) (err error) { rhl.Printf("curl-based fetcher called for '%s'", ctx.SourceUri) @@ -39,13 +89,19 @@ func FetchFileCurl(ctx *ImportContext, uri *url.URL) (err error) { defer easy.Cleanup() if easy != nil { - ctx.SourceFile = ctx.Config.TempDir + "/" + path.Base(uri.Path) + easy.Setopt(curl.OPT_FOLLOWLOCATION, true) easy.Setopt(curl.OPT_URL, ctx.SourceUri) - easy.Setopt(curl.OPT_WRITEFUNCTION, func(ptr []byte, userdata interface{}) bool { - // TODO: actually store data to ctx.SourceFile - return true - }) + + cbdata := &CurlCBData{remotename: path.Base(uri.Path)} + defer cbdata.Cleanup() + cbdata.basepath = ctx.Config.TempDir // TODO: create temporary directory + + easy.Setopt(curl.OPT_HEADERFUNCTION, curlHeaderCallback) + easy.Setopt(curl.OPT_HEADERDATA, cbdata) + + easy.Setopt(curl.OPT_WRITEFUNCTION, curlWriteCallback) + easy.Setopt(curl.OPT_WRITEDATA, cbdata) easy.Setopt(curl.OPT_NOPROGRESS, false) started := int64(0) @@ -56,12 +112,14 @@ func FetchFileCurl(ctx *ImportContext, uri *url.URL) (err error) { fmt.Printf("Downloaded: %3.2f%%, Speed: %.1fKiB/s \r", dlnow/dltotal*100, dlnow/1000/float64((time.Now().Unix()-started))) return true }) + easy.Setopt(curl.OPT_PROGRESSDATA, ctx) if err = easy.Perform(); err != nil { return } fmt.Printf("\n") + ctx.SourceFile = cbdata.filename ctx.DeleteSourceFile = true } return |