add custom url extractor

CHANGELOG.md

@@ -5,14 +5,15 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/)
 and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
-## Added
+### Added
 - Clearer debug
-- Flags (CLI options): debug, error video file, port, version
-## Removed
+- Flags (CLI options): debug, error video file, port, version, extractor select
+- Custom url extractor
+### Removed
 - Port passing as argument
 
 ## 0.4.0 - 2020-08-03
-- technical update
+Technical update
 
 ## 0.3.0 - 2017-11-03
 ### Added

parser.go

@@ -40,43 +40,57 @@ type corruptedT struct {
 	size int64
 }
 
-func youtubeDL(vURL, vHeight, vFormat string, debug debugT) (string, int64, error) {
-	// videoFormat = "(mp4)[height<=720]"
-	videoFormat := "(" + vFormat + ")[height<=" + vHeight + "]"
-	cmd := exec.Command("youtube-dl", "-f", videoFormat, "-g", vURL)
-	out, err := runCmd(cmd)
-	if err != nil {
-		return "", 0, err
-	}
-	output := strings.TrimSpace(string(out))
-	var expire int64
-	u, err := url.Parse(output)
-	if err == nil {
-		m, _ := url.ParseQuery(u.RawQuery)
-		if e, ok := m["expire"]; ok {
-			e1, err1 := strconv.ParseInt(e[0], 10, 64)
-			if err1 == nil {
-				expire = e1
+type extractorF func(string, string, string, debugT) (string, int64, error)
+
+func getYTDL() extractorF {
+	return func(vURL, vHeight, vFormat string, debug debugT) (string, int64, error) {
+		// videoFormat = "(mp4)[height<=720]"
+		videoFormat := "(" + vFormat + ")[height<=" + vHeight + "]"
+		cmd := exec.Command("youtube-dl", "-f", videoFormat, "-g", vURL)
+		out, err := runCmd(cmd)
+		if err != nil {
+			return "", 0, err
+		}
+		var expire int64
+		u, err := url.Parse(out)
+		if err == nil {
+			m, _ := url.ParseQuery(u.RawQuery)
+			if e, ok := m["expire"]; ok {
+				e1, err1 := strconv.ParseInt(e[0], 10, 64)
+				if err1 == nil {
+					expire = e1
+				}
 			}
 		}
+		return out, expire, err
+	}
+}
+
+func getCustomDL(path string) extractorF {
+	return func(vURL, vHeight, vFormat string, debug debugT) (string, int64, error) {
+		cmd := exec.Command(path, vURL, vHeight, vFormat)
+		out, err := runCmd(cmd)
+		return out, 0, err
 	}
-	return output, expire, err
 }
 
 func runCmd(cmd *exec.Cmd) (string, error) {
 	var stdout, stderr bytes.Buffer
 	cmd.Stdout = &stdout
 	cmd.Stderr = &stderr
 	err := cmd.Run()
-	outStr, errStr := string(stdout.Bytes()), string(stderr.Bytes())
+	toS := func(s bytes.Buffer) string {
+		return strings.TrimSpace(string(s.Bytes()))
+	}
+	outStr, errStr := toS(stdout), toS(stderr)
 	if err != nil {
 		combinedErrStr := fmt.Sprintf("%s\n%s\n%s", err.Error(), outStr, errStr)
 		return "", errors.New(combinedErrStr)
 	}
-	return outStr, nil
+	return strings.TrimSpace(outStr), nil
 }
 
-func getLink(vidurl string, debug debugT, links *linksCache) (string, error) {
+func getLink(vidurl string, debug debugT, links *linksCache, extractor extractorF) (string, error) {
 	now := time.Now().Unix()
 	vidurl = strings.TrimSpace(vidurl)
 	splitted := strings.Split(vidurl, "?/?")
@@ -116,7 +130,7 @@ func getLink(vidurl string, debug debugT, links *linksCache) (string, error) {
 	if ok {
 		return lnk.url, nil
 	}
-	url, expire, err := youtubeDL(vidurl, vh, vf, debug)
+	url, expire, err := extractor(vidurl, vh, vf, debug)
 	if expire == 0 {
 		expire = now + 10800
 	}

streamer.go

@@ -20,21 +20,31 @@ func main() {
 	var enableDebug bool
 	var portInt uint
 	var errorVideoPath string
+	var youtubedl bool
+	var customdl string
 	flag.BoolVar(&version, "version", false, "prints current yt-proxy version")
 	flag.BoolVar(&enableDebug, "debug", false, "turn on debug")
 	flag.UintVar(&portInt, "port", 8080, "listen port")
 	flag.StringVar(&errorVideoPath, "error-video", "corrupted.mp4", "file that will be shown on errors")
+	flag.BoolVar(&youtubedl, "youtube-dl", true, "use youtube-dl as url extractor")
+	flag.StringVar(&customdl, "custom-extractor", "", "use custom url extractor")
 	flag.Parse()
 	if version {
 		fmt.Println(appVersion)
 		os.Exit(0)
 	}
+	var extractor extractorF
+	if len(customdl) > 0 {
+		extractor = getCustomDL(customdl)
+	} else {
+		extractor = getYTDL()
+	}
 	var requests chan requestChan
 	requests = make(chan requestChan)
 	var links linksCache
 	links.cache = make(map[string]lnkT)
 	debug := getDebugFunc(enableDebug)
-	go parseLinks(requests, debug, &links)
+	go parseLinks(requests, debug, &links, extractor)
 	errorVideo := readErrorVideo(errorVideoPath)
 	port := fmt.Sprintf("%d", portInt)
 	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
@@ -108,7 +118,7 @@ func playVideo(w http.ResponseWriter, req *http.Request, requests chan requestCh
 			}
 		}
 	} else {
-		log.Println("yotube-dl error:", r.err)
+		log.Println("URL extractor error:", r.err)
 	}
 
 	if success == false {
@@ -119,11 +129,11 @@ func playVideo(w http.ResponseWriter, req *http.Request, requests chan requestCh
 	log.Printf("%s disconnected\n", req.RemoteAddr)
 }
 
-func parseLinks(requests <-chan requestChan, debug debugT, links *linksCache) {
+func parseLinks(requests <-chan requestChan, debug debugT, links *linksCache, extractor extractorF) {
 	for {
 		r := <-requests
 		url := r.url
-		rURL, rErr := getLink(url, debug, links)
+		rURL, rErr := getLink(url, debug, links, extractor)
 		debug("Extractor returned URL", rURL)
 		debug("Extractor returned error", rErr)
 		r.answerChan <- response{url: rURL, err: rErr}