MacFastLookup/service/download.go

84 lines
1.8 KiB
Go
Raw Normal View History

2024-07-23 16:32:49 +08:00
package service
import (
"fmt"
"io"
"net/http"
"os"
"strconv"
)
// ProgressWriter is a custom writer to track progress and execute a callback
type ProgressWriter struct {
io.Writer
Total int64
Progress int64
Callback func(progress int64, total int64)
}
func (pw *ProgressWriter) Write(p []byte) (int, error) {
n, err := pw.Writer.Write(p)
if err != nil {
return n, err
}
pw.Progress += int64(n)
pw.Callback(pw.Progress, pw.Total)
return n, nil
}
func DownloadFile(url string, filePath string, overwrite bool, callback func(progress int64, total int64)) error {
// Check if the file exists
if _, err := os.Stat(filePath); err == nil {
if overwrite {
// Overwrite the file if it exists
if err := os.Remove(filePath); err != nil {
return err
}
} else {
return fmt.Errorf("file %s already exists", filePath)
}
}
// Create the file
out, err := os.Create(filePath)
if err != nil {
return err
}
defer out.Close()
// Get the data
resp, err := http.Get(url)
if err != nil {
return err
}
defer resp.Body.Close()
// Check server response
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("bad status: %s", resp.Status)
}
// Get the size
size, err := strconv.Atoi(resp.Header.Get("Content-Length"))
if err != nil {
return err
}
// Create progress writer
pw := &ProgressWriter{
Writer: out,
Total: int64(size),
Callback: callback,
}
// Write the body to file with progress
_, err = io.Copy(pw, resp.Body)
if err != nil {
return err
}
fmt.Println("\nDownload complete")
return nil
}