initial Go port of streamrip

internal/download/downloader.go

@@ -0,0 +1,200 @@
+package download
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"sync/atomic"
+	"time"
+
+	"github.com/vbauerster/mpb/v8"
+	"github.com/vbauerster/mpb/v8/decor"
+	"golang.org/x/term"
+
+	"streamrip-go/internal/netutil"
+)
+
+type Downloader struct {
+	http         *http.Client
+	showProgress bool
+	progress     *mpb.Progress
+	barStarted   atomic.Int32
+}
+
+func New() *Downloader {
+	return NewWithOptions(true, true)
+}
+
+func NewWithVerifySSL(verifySSL bool) *Downloader {
+	return NewWithOptions(verifySSL, true)
+}
+
+func NewWithOptions(verifySSL bool, showProgress bool) *Downloader {
+	forceProgress := strings.EqualFold(os.Getenv("STREAMRIP_GO_FORCE_PROGRESS"), "1") || strings.EqualFold(os.Getenv("STREAMRIP_GO_FORCE_PROGRESS"), "true")
+	interactive := showProgress && (forceProgress || (term.IsTerminal(int(os.Stderr.Fd())) && strings.ToLower(os.Getenv("TERM")) != "dumb"))
+	d := &Downloader{http: netutil.NewHTTPClient(2*time.Minute, verifySSL), showProgress: interactive}
+	if interactive {
+		d.progress = mpb.New(mpb.WithWidth(40), mpb.WithOutput(os.Stderr))
+	}
+	return d
+}
+
+func (d *Downloader) File(ctx context.Context, sourceURL, outputPath string) error {
+	return d.file(ctx, sourceURL, outputPath, true)
+}
+
+func (d *Downloader) FileNoProgress(ctx context.Context, sourceURL, outputPath string) error {
+	return d.file(ctx, sourceURL, outputPath, false)
+}
+
+func (d *Downloader) file(ctx context.Context, sourceURL, outputPath string, allowProgress bool) error {
+	if err := os.MkdirAll(filepath.Dir(outputPath), 0o755); err != nil {
+		return err
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, sourceURL, nil)
+	if err != nil {
+		return err
+	}
+
+	resp, err := d.http.Do(req)
+	if err != nil {
+		return err
+	}
+	defer func() { _ = resp.Body.Close() }()
+
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("download failed: status=%d", resp.StatusCode)
+	}
+
+	reader := bufio.NewReader(resp.Body)
+	peek, _ := reader.Peek(1024)
+	if isManifestResponse(resp.Header.Get("Content-Type"), peek) {
+		_ = resp.Body.Close()
+		return d.streamManifestWithFFmpeg(ctx, sourceURL, outputPath)
+	}
+
+	out, err := os.Create(outputPath)
+	if err != nil {
+		return err
+	}
+	defer func() { _ = out.Close() }()
+
+	if d.ProgressEnabled() && allowProgress && resp.ContentLength > 0 {
+		d.barStarted.Store(1)
+		desc := shortenName(filepath.Base(outputPath), 54)
+		bar := d.progress.AddBar(
+			resp.ContentLength,
+			mpb.PrependDecorators(
+				decor.Name(desc+" ", decor.WC{W: 56, C: decor.DSyncWidth | decor.DindentRight}),
+				decor.Percentage(decor.WCSyncWidthR),
+			),
+			mpb.AppendDecorators(
+				decor.CountersKibiByte("% .1f / % .1f", decor.WCSyncWidthR),
+				decor.Name(" | ", decor.WCSyncWidth),
+				decor.AverageSpeed(decor.SizeB1024(0), "% .1f", decor.WCSyncWidthR),
+				decor.Name(" | ETA ", decor.WCSyncWidth),
+				decor.AverageETA(decor.ET_STYLE_GO, decor.WCSyncWidthR),
+			),
+			mpb.BarRemoveOnComplete(),
+		)
+		buf := make([]byte, 256*1024)
+		for {
+			n, readErr := reader.Read(buf)
+			if n > 0 {
+				if _, writeErr := out.Write(buf[:n]); writeErr != nil {
+					return writeErr
+				}
+				bar.IncrBy(n)
+			}
+			if readErr != nil {
+				if readErr == io.EOF {
+					break
+				}
+				return readErr
+			}
+		}
+	} else {
+		if _, err = io.Copy(out, reader); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (d *Downloader) Close() {
+	if d.progress != nil {
+		d.progress.Wait()
+	}
+}
+
+func (d *Downloader) ProgressEnabled() bool {
+	return d.showProgress && d.progress != nil
+}
+
+func (d *Downloader) Logf(format string, args ...any) {
+	msg := fmt.Sprintf(format, args...)
+	if d.ProgressEnabled() && d.barStarted.Load() == 1 {
+		_, _ = d.progress.Write([]byte(msg))
+		return
+	}
+	fmt.Print(msg)
+}
+
+func shortenName(name string, max int) string {
+	if max <= 0 {
+		return name
+	}
+	r := []rune(name)
+	if len(r) <= max {
+		return name
+	}
+	if max <= 3 {
+		return string(r[:max])
+	}
+	return string(r[:max-3]) + "..."
+}
+
+func (d *Downloader) streamManifestWithFFmpeg(ctx context.Context, sourceURL, outputPath string) error {
+	if _, err := exec.LookPath("ffmpeg"); err != nil {
+		return fmt.Errorf("ffmpeg not found for manifest stream: %w", err)
+	}
+
+	args := []string{
+		"-y",
+		"-protocol_whitelist", "file,http,https,tcp,tls,crypto,data",
+		"-i", sourceURL,
+		"-map", "0:a:0",
+		"-c", "copy",
+		outputPath,
+	}
+
+	cmd := exec.CommandContext(ctx, "ffmpeg", args...)
+	output, err := cmd.CombinedOutput()
+	if err != nil {
+		return fmt.Errorf("ffmpeg stream copy failed: %w: %s", err, string(output))
+	}
+	return nil
+}
+
+func isManifestResponse(contentType string, peek []byte) bool {
+	ct := strings.ToLower(contentType)
+	if strings.Contains(ct, "dash+xml") || strings.Contains(ct, "mpegurl") || strings.Contains(ct, "vnd.apple.mpegurl") {
+		return true
+	}
+	s := strings.TrimSpace(strings.ToLower(string(peek)))
+	if strings.HasPrefix(s, "<?xml") && strings.Contains(s, "<mpd") {
+		return true
+	}
+	if strings.HasPrefix(s, "#extm3u") {
+		return true
+	}
+	return false
+}