Merge pull request #40 from sonroyaalmerol/user-agent-stream

Use USER_AGENT env var for the actual video streams

m3u/parser.go

@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"io"
 	"log"
-	"net/http"
 	"os"
 	"regexp"
 	"strconv"
@@ -16,6 +15,7 @@ import (
 	"time"
 
 	"m3u-stream-merger/database"
+	"m3u-stream-merger/utils"
 )
 
 func parseLine(line string, nextLine string, m3uIndex int) database.StreamInfo {
@@ -120,24 +120,9 @@ func insertStreamToDb(db *sql.DB, currentStream database.StreamInfo) error {
 }
 
 func downloadM3UToBuffer(m3uURL string, buffer *bytes.Buffer) (err error) {
-	// Set the custom User-Agent header
-	userAgent, userAgentExists := os.LookupEnv("USER_AGENT")
-	if !userAgentExists {
-		userAgent = "IPTV Smarters/1.0.3 (iPad; iOS 16.6.1; Scale/2.00)"
-	}
-
-	// Create a new HTTP client with a custom User-Agent header
-	client := &http.Client{
-		CheckRedirect: func(req *http.Request, via []*http.Request) error {
-			// Follow redirects while preserving the custom User-Agent header
-			req.Header.Set("User-Agent", userAgent)
-			return nil
-		},
-	}
-
 	// Download M3U for processing
 	log.Printf("Downloading M3U from URL: %s\n", m3uURL)
-	resp, err := client.Get(m3uURL)
+	resp, err := utils.CustomHttpRequest("GET", m3uURL)
 	if err != nil {
 		return fmt.Errorf("HTTP GET error: %v", err)
 	}

mp4_handler.go

@@ -39,7 +39,7 @@ func loadBalancer(stream database.StreamInfo) (resp *http.Response, selectedUrl
 				continue // Skip this stream if concurrency limit reached
 			}
 
-			resp, err = http.Get(url.Content)
+			resp, err = utils.CustomHttpRequest("GET", url.Content)
 			if err == nil {
 				selectedUrl = &url
 				break
@@ -62,7 +62,7 @@ func loadBalancer(stream database.StreamInfo) (resp *http.Response, selectedUrl
 				continue // Skip this stream if concurrency limit reached
 			}
 
-			resp, err = http.Get(url.Content)
+			resp, err = utils.CustomHttpRequest("GET", url.Content)
 			if err == nil {
 				selectedUrl = &url
 				break
@@ -79,7 +79,7 @@ func loadBalancer(stream database.StreamInfo) (resp *http.Response, selectedUrl
 		log.Printf("All concurrency limits have been reached. Falling back to connection checking mode...\n")
 		// Connection check mode
 		for _, url := range stream.URLs {
-			resp, err = http.Get(url.Content)
+			resp, err = utils.CustomHttpRequest("GET", url.Content)
 			if err == nil {
 				selectedUrl = &url
 				break

utils/env.go

@@ -0,0 +1,19 @@
+package utils
+
+import (
+	"os"
+)
+
+func GetEnv(env string) string {
+	switch env {
+	case "USER_AGENT":
+		// Set the custom User-Agent header
+		userAgent, userAgentExists := os.LookupEnv("USER_AGENT")
+		if !userAgentExists {
+			userAgent = "IPTV Smarters/1.0.3 (iPad; iOS 16.6.1; Scale/2.00)"
+		}
+		return userAgent
+	default:
+		return ""
+	}
+}

utils/http.go

@@ -0,0 +1,30 @@
+package utils
+
+import "net/http"
+
+func CustomHttpRequest(method string, url string) (*http.Response, error) {
+	userAgent := GetEnv("USER_AGENT")
+
+	// Create a new HTTP client with a custom User-Agent header
+	client := &http.Client{
+		CheckRedirect: func(req *http.Request, via []*http.Request) error {
+			// Follow redirects while preserving the custom User-Agent header
+			req.Header.Set("User-Agent", userAgent)
+			return nil
+		},
+	}
+
+	req, err := http.NewRequest(method, url, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Set("User-Agent", userAgent)
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+
+	return resp, nil
+}