abzcoding / hget

package main

import (

	"bufio"

	"flag"

	"io"

	"os"

	"os/signal"

	"runtime"

	"strings"

	"syscall"

	"github.com/imkira/go-task"

)

var displayProgress = true

func main() {

	// var err error

	var proxy, filePath, bwLimit, resumeTask string

	conn := flag.Int("n", runtime.NumCPU(), "number of connections")

	skiptls := flag.Bool("skip-tls", false, "skip certificate verification for https")

	flag.StringVar(&proxy, "proxy", "", "proxy for downloading, e.g. -proxy '127.0.0.1:12345' for socks5 or -proxy 'http://proxy.com:8080' for http proxy")

	flag.StringVar(&filePath, "file", "", "path to a file that contains one URL per line")

	flag.StringVar(&bwLimit, "rate", "", "bandwidth limit during download, e.g. -rate 10kB or -rate 10MiB")

	flag.StringVar(&resumeTask, "resume", "", "resume download task with given task name (or URL)")

	probe := flag.String("probe", "", "probe URL for range and content-length without downloading")

	flag.Parse()

	args := flag.Args()

	// Probe diagnostics mode

	if *probe != "" {

		DebugProbe(*probe, *skiptls, proxy)

		return

	}

	// If the resume flag is provided, use that path (ignoring other arguments)

	if resumeTask != "" {

		state, err := Resume(resumeTask)

		FatalCheck(err)

		Execute(state.URL, state, *conn, *skiptls, proxy, bwLimit)

		return

	}

	// If no resume flag, then check for positional URL or file input

	if len(args) < 1 {

		if len(filePath) < 1 {

			Errorln("A URL or input file with URLs is required")

			usage()

			os.Exit(1)

		}

		// Create a serial group for processing multiple URLs in a file.

		g1 := task.NewSerialGroup()

		file, err := os.Open(filePath)

		if err != nil {

			FatalCheck(err)

		}

		defer file.Close()

		reader := bufio.NewReader(file)

		for {

			line, _, err := reader.ReadLine()

			if err == io.EOF {

				break

			}

			url := strings.TrimSpace(string(line))

			if url == "" || strings.HasPrefix(url, "#") {

				continue

			}

			// Add the download task for each URL

			g1.AddChild(downloadTask(url, nil, *conn, *skiptls, proxy, bwLimit))

		}

		g1.Run(nil)

		return

	}

	// Otherwise, if a URL is provided as positional argument, treat it as a new download.

	downloadURL := args[0]

	// Check if a folder already exists for the task and remove if necessary.

	if ExistDir(FolderOf(downloadURL)) {

		Warnf("Downloading task already exists, remove it first \n")

		err := os.RemoveAll(FolderOf(downloadURL))

		FatalCheck(err)

	}

	Execute(downloadURL, nil, *conn, *skiptls, proxy, bwLimit)

}

func downloadTask(url string, state *State, conn int, skiptls bool, proxy string, bwLimit string) task.Task {

	run := func(t task.Task, ctx task.Context) {

		Execute(url, state, conn, skiptls, proxy, bwLimit)

	}

	return task.NewTaskWithFunc(run)

}

// Execute configures the HTTPDownloader and uses it to download the target.

func Execute(url string, state *State, conn int, skiptls bool, proxy string, bwLimit string) {

	// Capture OS interrupt signals

	signalChan := make(chan os.Signal, 1)

	signal.Notify(signalChan,

		syscall.SIGHUP,

		syscall.SIGINT,

		syscall.SIGTERM,

		syscall.SIGQUIT)

	var files = make([]string, 0)

	var parts = make([]Part, 0)

	var isInterrupted = false

	doneChan := make(chan bool, conn)

	fileChan := make(chan string, conn)

	errorChan := make(chan error, 1)

	stateChan := make(chan Part, 1)

	interruptChan := make(chan bool, conn)

	var downloader *HTTPDownloader

	if state == nil {

		downloader = NewHTTPDownloader(url, conn, skiptls, proxy, bwLimit)

	} else {

		downloader = &HTTPDownloader{

			url:       state.URL,

			file:      TaskFromURL(state.URL),

			par:       int64(len(state.Parts)),

			parts:     state.Parts,

			resumable: true,

		}

	}

	go downloader.Do(doneChan, fileChan, errorChan, interruptChan, stateChan)

	for {

		select {

		case <-signalChan:

			// Signal all active download routines to interrupt.

			isInterrupted = true

			for i := 0; i < conn; i++ {

				interruptChan <- true

			}

		case file := <-fileChan:

			files = append(files, file)

		case err := <-errorChan:

			Errorf("%v", err)

			panic(err)

		case part := <-stateChan:

			parts = append(parts, part)

		case <-doneChan:

			// Ensure we drain remaining part notifications before finalizing.

			numParts := len(downloader.parts)

			for len(files) < numParts {

				file := <-fileChan

				files = append(files, file)

			}

			for len(parts) < numParts {

				part := <-stateChan

				parts = append(parts, part)

			}

			if isInterrupted {

				if downloader.resumable {