youtube live to m3u

Signed-off-by: jet tsang zeon-git <zeon-git@jettsang.com>

1 files changed, 48 insertions, 11 deletions

diff --git a/main.go b/main.go
index fc456d4..66f5d6d 100644
--- a/main.go
+++ b/main.go
@@ -25,6 +25,7 @@ var shortenMap map[string]string
 var listen string
 var reloadURI = "/reload"
 var clientHeaderURI = "/header"
+var UserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/103.0.5060.53 Safari/537.36"
 
 func init() {
 	const (
@@ -43,40 +44,63 @@ func init() {
 
 	flag.StringVar(&listen, "listen", defaultListen, listenUsage)
 	flag.StringVar(&listen, "l", defaultListen, listenUsage+" (shorthand)")
-
 	shortenMap = map[string]string{}
 }
 func main() {
 	flag.Parse()
 	go autoLoad()
 	log.Println("Listening to:", listen)
-	if err := fasthttp.ListenAndServe(listen, requestHandler); err != nil {
+	s := &fasthttp.Server{
+		Handler:               requestHandler,
+		NoDefaultServerHeader: true,
+	}
+
+	if err := s.ListenAndServe(listen); err != nil {
 		log.Fatalf("Error in ListenAndServe: %s", err)
 	}
 }
 
 func requestHandler(ctx *fasthttp.RequestCtx) {
-	if string(ctx.RequestURI()) == reloadURI {
+	strr := string(ctx.RequestURI())
+	ctx.Response.Header.Add("Server", "jet's relay server")
+	if strr == reloadURI {
 		loadData()
 		fmt.Fprint(ctx, "ok")
 		ctx.SetStatusCode(fasthttp.StatusAccepted)
 		return
-	} else if string(ctx.RequestURI()) == clientHeaderURI {
+	} else if strr == clientHeaderURI {
 		fmt.Fprint(ctx, string(ctx.Request.Header.Header()))
 		ctx.SetStatusCode(fasthttp.StatusAccepted)
 		return
+	} else if strings.HasPrefix(strr, "/yt/") {
+		var req fasthttp.Request
+		var resp fasthttp.Response
+		ythash := strings.TrimLeft(strr, "/yt/")
+		req.Header.Add("User-Agent", UserAgent)
+		req.Header.Add("Referer", "https://www.youtube.com")
+		req.SetRequestURI("https://www.youtube.com/watch?v=" + ythash)
+		fasthttp.Do(&req, &resp)
+		ytbody := string(resp.Body())
+		m3u8 := GetBetweenStr(ytbody, "hlsManifestUrl\":\"", "\"},\"heart")
+		ctx.Redirect(m3u8, 302)
+		defer func() {
+			fasthttp.ReleaseRequest(&req)
+			req.ConnectionClose()
+			resp.ConnectionClose()
+		}()
+		return
 	}
 	if relayUrl, exist := shortenMap[string(ctx.Path())]; exist {
-		var newRequest fasthttp.Request
-		ctx.Request.Header.CopyTo(&newRequest.Header)
-		newRequest.SetRequestURI(relayUrl)
-		newRequest.Header.SetBytesV("Host", newRequest.Host())
+		var req fasthttp.Request
+		ctx.Request.Header.CopyTo(&req.Header)
+		req.SetRequestURI(relayUrl)
+		req.Header.SetBytesV("Host", req.Host())
 		ctx.URI().QueryArgs().VisitAll(func(key []byte, value []byte) {
-			newRequest.URI().QueryArgs().AddBytesKV(key, value)
+			req.URI().QueryArgs().AddBytesKV(key, value)
 		})
-		defer fasthttp.ReleaseRequest(&newRequest)
+		defer fasthttp.ReleaseRequest(&req)
 
-		fasthttp.Do(&newRequest, &ctx.Response)
+		fasthttp.Do(&req, &ctx.Response)
 	} else {
 		ctx.Error("err", 404)
 	}
@@ -116,3 +140,16 @@ func loadData() {
 	log.Println("urls mapped.")
 
 }
+func GetBetweenStr(str, start, end string) string {
+	n := strings.Index(str, start) + len(start)
+	if n == -1 {
+		n = 0
+	}
+	str = string([]byte(str)[n:])
+	m := strings.Index(str, end)
+	if m == -1 {
+		m = len(str)
+	}
+	str = string([]byte(str)[:m])
+	return str
+}