Skip to content

Commit 82f4f77

Browse files
authored
Feat/yt dlp (#8)
* trying to add youtube download support Signed-off-by: abzcoding <abzcoding@gmail.com> * support cookies+youtube Signed-off-by: abzcoding <abzcoding@gmail.com> * cleanup Signed-off-by: abzcoding <abzcoding@gmail.com> --------- Signed-off-by: abzcoding <abzcoding@gmail.com>
1 parent 926e397 commit 82f4f77

17 files changed

Lines changed: 6087 additions & 1 deletion

cmd/hget/main.go

Lines changed: 228 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,21 @@
11
package main
22

33
import (
4+
"bufio"
45
"context"
56
"errors"
67
"flag"
78
"fmt"
89
"os"
910
"os/signal"
1011
"runtime"
12+
"strings"
1113
"syscall"
1214
"time"
1315

1416
"github.com/abzcoding/hget/internal/batch"
1517
"github.com/abzcoding/hget/internal/downloader"
18+
"github.com/abzcoding/hget/internal/extractor"
1619
"github.com/abzcoding/hget/internal/state"
1720
"github.com/abzcoding/hget/internal/ui"
1821
"github.com/abzcoding/hget/internal/util"
@@ -24,7 +27,9 @@ var GitCommit string
2427
func main() {
2528
flag.Usage = ui.PrintHelp
2629

27-
var proxy, filePath, bwLimit, resumeTask string
30+
var proxy, filePath, bwLimit, resumeTask, extractorMode, cookiesFile, cookiesBrowser string
31+
var quality, container, lang string
32+
var pickFormat bool
2833

2934
conn := flag.Int("n", runtime.NumCPU(), "number of connections")
3035
skiptls := flag.Bool("skip-tls", false, "skip certificate verification for https")
@@ -33,6 +38,13 @@ func main() {
3338
flag.StringVar(&filePath, "file", "", "path to a file that contains one URL per line")
3439
flag.StringVar(&bwLimit, "rate", "", "bandwidth limit during download, e.g. -rate 10kB or -rate 10MiB")
3540
flag.StringVar(&resumeTask, "resume", "", "resume download task with given task name (or URL)")
41+
flag.StringVar(&extractorMode, "extractor", "auto", "extractor mode: auto | yt-dlp | none (auto picks yt-dlp for known media hosts)")
42+
flag.StringVar(&cookiesFile, "cookies", "", "path to Netscape-format cookies.txt for the extractor (forwarded to yt-dlp --cookies)")
43+
flag.StringVar(&cookiesBrowser, "cookies-from-browser", "", "browser to extract cookies from for the extractor, e.g. firefox, chrome:Default (forwarded to yt-dlp --cookies-from-browser)")
44+
flag.StringVar(&quality, "quality", "720p", "extractor quality preset: 360p | 480p | 720p | 1080p | 1440p | 4K | 8K | best | audio")
45+
flag.StringVar(&container, "container", "mp4", "extractor output container: mp4 | mkv | webm")
46+
flag.StringVar(&lang, "audio-lang", "en", "preferred audio language for the extractor (forwarded as yt-dlp -S lang:<code>); empty disables the bias")
47+
flag.BoolVar(&pickFormat, "pick-format", false, "open the VCR rocker UI to pick resolution/audio/container by hand instead of using --quality")
3648
probe := flag.String("probe", "", "probe URL for range and content-length without downloading")
3749
timeout := flag.Duration("timeout", 15*time.Second, "timeout for awaiting response headers (e.g., 30s, 1m)")
3850

@@ -71,12 +83,40 @@ func main() {
7183
ui.PrintHelp()
7284
os.Exit(1)
7385
}
86+
// Route to the yt-dlp shelf when any URL in the file looks
87+
// extractable (or --extractor=yt-dlp forces it). All-or-
88+
// nothing: the whole list is treated as a video batch and any
89+
// plain HTTP URL falls through yt-dlp's generic extractor.
90+
if urls, useExtractor, err := loadBatchURLs(filePath, extractorMode); err != nil {
91+
ui.ShowMessage(ui.MessageError, "FILE ERROR", err.Error())
92+
os.Exit(1)
93+
} else if useExtractor {
94+
runExtractorBatch(rootCtx, urls, extractor.Options{
95+
CookiesFile: cookiesFile,
96+
CookiesFromBrowser: cookiesBrowser,
97+
LangPref: lang,
98+
}, extractor.QualityPreset(quality, container), pickFormat)
99+
return
100+
}
74101
batch.RunBatchDownloads(rootCtx, filePath, *conn, *skiptls, proxy, bwLimit, *timeout, *verify)
75102
return
76103
}
77104

78105
// Single URL download.
79106
downloadURL := args[0]
107+
108+
// Extractor mode (yt-dlp pipeline). Picked when explicitly forced
109+
// or when --extractor=auto and the URL host matches a known media
110+
// site that hget's HTTP engine can't handle directly (YouTube etc.).
111+
if shouldUseExtractor(extractorMode, downloadURL) {
112+
runExtractor(rootCtx, downloadURL, extractor.Options{
113+
CookiesFile: cookiesFile,
114+
CookiesFromBrowser: cookiesBrowser,
115+
LangPref: lang,
116+
}, extractor.QualityPreset(quality, container), pickFormat)
117+
return
118+
}
119+
80120
destFile := util.TaskFromURL(downloadURL)
81121

82122
// Check if final file already exists
@@ -166,6 +206,193 @@ func runResume(rootCtx context.Context, resumeTask string, conn int, skiptls boo
166206
runOne(rootCtx, st.URL, st, conn, skiptls, proxy, bwLimit, timeout)
167207
}
168208

209+
// shouldUseExtractor decides whether the URL should be routed through
210+
// the yt-dlp pipeline. Modes:
211+
// - "yt-dlp": always use yt-dlp
212+
// - "auto": use yt-dlp when the host looks like a media site
213+
// - "none": never use yt-dlp (force the plain HTTP engine)
214+
func shouldUseExtractor(mode, url string) bool {
215+
switch mode {
216+
case "yt-dlp", "ytdlp":
217+
return true
218+
case "none", "off", "false":
219+
return false
220+
default: // "auto" and unknown values fall through to detection
221+
return extractor.LooksExtractable(url)
222+
}
223+
}
224+
225+
// runExtractor drives the yt-dlp pipeline behind the VCR + Mixer TUI.
226+
// On success the resolved output file is left in the current working
227+
// directory (yt-dlp's default), matching hget's existing behaviour.
228+
//
229+
// `preset` is the quality+container chosen via --quality / --container.
230+
// When `pickFormat` is false (the default), the rocker UI never appears
231+
// and the preset is fed straight into yt-dlp. When true, the user
232+
// gets a chance to override on a per-tape basis via the VCR's rockers.
233+
//
234+
// Cookie sources are validated BEFORE we start the TUI so a typo'd path
235+
// produces a clean error in the terminal instead of a cryptic message
236+
// flashing inside the alt-screen for half a second before exit.
237+
func runExtractor(rootCtx context.Context, url string, opts extractor.Options, preset extractor.FormatSelection, pickFormat bool) {
238+
if opts.CookiesFile != "" {
239+
if _, err := os.Stat(opts.CookiesFile); err != nil {
240+
ui.ShowMessage(ui.MessageError, "COOKIES FILE NOT FOUND",
241+
fmt.Sprintf("--cookies %s: %v", opts.CookiesFile, err))
242+
os.Exit(1)
243+
}
244+
}
245+
if opts.CookiesFile != "" && opts.CookiesFromBrowser != "" {
246+
ui.ShowMessage(ui.MessageWarning, "COOKIE SOURCES",
247+
"both --cookies and --cookies-from-browser are set; yt-dlp will pick the browser source")
248+
}
249+
250+
itemCtx, cancelItem := context.WithCancelCause(rootCtx)
251+
defer cancelItem(nil)
252+
253+
err := ui.RunExtractorTUI(ui.ExtractorRunOptions{
254+
Ctx: itemCtx,
255+
URL: url,
256+
OnQuit: func() { cancelItem(downloader.ErrUserQuit) },
257+
}, func(sel ui.ExtractorSelector) error {
258+
picker := buildPicker(preset, pickFormat, sel)
259+
return extractor.Pipeline(itemCtx, url, "", opts, pickFormat, picker)
260+
})
261+
262+
if err != nil &&
263+
!errors.Is(err, downloader.ErrUserQuit) &&
264+
!errors.Is(err, downloader.ErrAbortBatch) &&
265+
!errors.Is(err, context.Canceled) {
266+
ui.Errorln(err)
267+
os.Exit(1)
268+
}
269+
}
270+
271+
// buildPicker constructs the SelectorFunc handed to the extractor
272+
// pipeline. Two flavours:
273+
//
274+
// - pickFormat = false (default): a fast-path picker that returns
275+
// the CLI preset immediately, never touching the UI. The VCR
276+
// stays in standby until yt-dlp starts streaming bytes.
277+
//
278+
// - pickFormat = true: blocks on the TUI's REC commit so the user
279+
// can manipulate the rockers. The selection's adaptive
280+
// descriptors propagate to subsequent tapes via the batch
281+
// FormatAll policy.
282+
func buildPicker(preset extractor.FormatSelection, pickFormat bool, sel ui.ExtractorSelector) extractor.SelectorFunc {
283+
if !pickFormat {
284+
return func(ctx context.Context, _ extractor.Meta) (extractor.FormatSelection, error) {
285+
return preset, nil
286+
}
287+
}
288+
return func(ctx context.Context, _ extractor.Meta) (extractor.FormatSelection, error) {
289+
s, err := sel(ctx)
290+
if err != nil {
291+
return extractor.FormatSelection{}, err
292+
}
293+
return extractor.FormatSelection{
294+
Spec: s.Spec,
295+
Container: s.Container,
296+
Pref: extractor.FormatPreference{
297+
HeightCeiling: s.HeightCeiling,
298+
FPSFloor: s.FPSFloor,
299+
VCodec: s.VCodec,
300+
ABRCeiling: s.ABRCeiling,
301+
Progressive: s.Progressive,
302+
},
303+
}, nil
304+
}
305+
}
306+
307+
// loadBatchURLs reads the URL list and decides which pipeline owns it.
308+
// The rule is all-or-nothing: as soon as one URL looks extractable (or
309+
// --extractor=yt-dlp forces it) the whole list goes to yt-dlp. Comment
310+
// lines (#) and blanks are stripped to match batch.RunBatchDownloads.
311+
func loadBatchURLs(filePath, extractorMode string) ([]string, bool, error) {
312+
f, err := os.Open(filePath)
313+
if err != nil {
314+
return nil, false, fmt.Errorf("could not open URL list: %s\n%v", filePath, err)
315+
}
316+
defer f.Close()
317+
var urls []string
318+
scanner := bufio.NewScanner(f)
319+
for scanner.Scan() {
320+
line := strings.TrimSpace(scanner.Text())
321+
if line == "" || strings.HasPrefix(line, "#") {
322+
continue
323+
}
324+
urls = append(urls, line)
325+
}
326+
if scanErr := scanner.Err(); scanErr != nil {
327+
return nil, false, fmt.Errorf("error reading file: %s\n%v", filePath, scanErr)
328+
}
329+
if len(urls) == 0 {
330+
return nil, false, fmt.Errorf("no URLs found in: %s", filePath)
331+
}
332+
// Forced modes short-circuit the per-URL detection.
333+
switch extractorMode {
334+
case "yt-dlp", "ytdlp":
335+
return urls, true, nil
336+
case "none", "off", "false":
337+
return urls, false, nil
338+
}
339+
for _, u := range urls {
340+
if extractor.LooksExtractable(u) {
341+
return urls, true, nil
342+
}
343+
}
344+
return urls, false, nil
345+
}
346+
347+
// runExtractorBatch drives the yt-dlp pipeline over a list of URLs
348+
// behind a single persistent VCR + cassette shelf TUI. All cookies
349+
// validation, signal handling, and exit semantics mirror runExtractor.
350+
func runExtractorBatch(rootCtx context.Context, urls []string, opts extractor.Options, preset extractor.FormatSelection, pickFormat bool) {
351+
if opts.CookiesFile != "" {
352+
if _, err := os.Stat(opts.CookiesFile); err != nil {
353+
ui.ShowMessage(ui.MessageError, "COOKIES FILE NOT FOUND",
354+
fmt.Sprintf("--cookies %s: %v", opts.CookiesFile, err))
355+
os.Exit(1)
356+
}
357+
}
358+
if opts.CookiesFile != "" && opts.CookiesFromBrowser != "" {
359+
ui.ShowMessage(ui.MessageWarning, "COOKIE SOURCES",
360+
"both --cookies and --cookies-from-browser are set; yt-dlp will pick the browser source")
361+
}
362+
363+
itemCtx, cancelItem := context.WithCancelCause(rootCtx)
364+
defer cancelItem(nil)
365+
366+
err := ui.RunExtractorTUI(ui.ExtractorRunOptions{
367+
Ctx: itemCtx,
368+
// The shelf shows the queue; the source line above the deck
369+
// gets per-tape updates via ExtractorURLMsg. Seed with the
370+
// first URL so the very first frame has something legible.
371+
URL: urls[0],
372+
OnQuit: func() { cancelItem(downloader.ErrUserQuit) },
373+
}, func(sel ui.ExtractorSelector) error {
374+
// One persistent selector — the batch policy decides whether
375+
// to actually invoke it for each tape. BatchFormatAll is the
376+
// default: the user picks once on tape #1; the adaptive
377+
// descriptors travel with the cached selection so subsequent
378+
// tapes that lack the original format IDs still resolve to a
379+
// close match (yt-dlp filter syntax with progressive fallback).
380+
//
381+
// When `pickFormat` is false the picker collapses to a fast
382+
// preset-returner; the rocker UI never shows.
383+
picker := buildPicker(preset, pickFormat, sel)
384+
return extractor.BatchPipeline(itemCtx, urls, "", opts, extractor.BatchFormatAll, pickFormat, picker)
385+
})
386+
387+
if err != nil &&
388+
!errors.Is(err, downloader.ErrUserQuit) &&
389+
!errors.Is(err, downloader.ErrAbortBatch) &&
390+
!errors.Is(err, context.Canceled) {
391+
ui.Errorln(err)
392+
os.Exit(1)
393+
}
394+
}
395+
169396
func runOne(rootCtx context.Context, url string, st *state.State, conn int, skiptls bool, proxy, bwLimit string, timeout time.Duration) {
170397
itemCtx, cancelItem := context.WithCancelCause(rootCtx)
171398
defer cancelItem(nil)

0 commit comments

Comments
 (0)