Enhance song downloading functionality: add support for YouTube cookies and update related methods

Author: mervyn-teo

internal/music/music.go

@@ -8,6 +8,7 @@ import (
 	"github.com/kkdai/youtube/v2"
 	"io"
 	"log"
+	"net/http"
 	"os"
 	"path/filepath"
 	"sync"
@@ -30,14 +31,28 @@ type Song struct {
 }
 
 // DownloadSong downloads the song from the given URL and saves it to a file. returns the file path to the downloaded song.
-func DownloadSong(url string) (filePath string, err error) {
+func DownloadSong(url string, ytbCookie string) (filePath string, err error) {
 	err = os.MkdirAll("./songCache", os.ModePerm)
 	if err != nil {
 		return "", fmt.Errorf("error creating directory: %w", err)
 	}
 
-	client := youtube.Client{}
-	video, err := client.GetVideo(url)
+	ytbClient := youtube.Client{}
+
+	if ytbCookie == "" {
+		return "", errors.New("no cookies supplied")
+	}
+
+	httpClient := http.Client{}
+	req, err := http.NewRequest("GET", "", nil)
+	if err != nil {
+		return "", fmt.Errorf("error creating request: %w", err)
+	}
+	req.Header.Set("Cookie", ytbCookie)
+
+	ytbClient.HTTPClient = &httpClient
+
+	video, err := ytbClient.GetVideo(url)
 
 	if err != nil {
 		return "", fmt.Errorf("error getting video: %w", err)
@@ -52,7 +67,7 @@ func DownloadSong(url string) (filePath string, err error) {
 	// Get the best audio format
 	audioFormat := video.Formats.WithAudioChannels()
 
-	stream, _, err := client.GetStream(video, &audioFormat[0])
+	stream, _, err := ytbClient.GetStream(video, &audioFormat[0])
 	if err != nil {
 		panic(err)
 	}
@@ -90,7 +105,7 @@ func DownloadSong(url string) (filePath string, err error) {
 }
 
 // PlaySong plays the audio file using the voice connection. gid and cid are the guild and channel IDs.
-func (s *SongList) PlaySong(gid string, cid string, bot *bot.Bot) error {
+func (s *SongList) PlaySong(gid string, cid string, bot *bot.Bot, ytbCookie string) error {
 	// check if the song list is empty
 	if len(s.Songs) <= 0 {
 		return fmt.Errorf("no songs in the list")
@@ -103,7 +118,7 @@ func (s *SongList) PlaySong(gid string, cid string, bot *bot.Bot) error {
 	}
 
 	currSong := s.Songs[0]
-	filePath, err := DownloadSong(currSong.Url)
+	filePath, err := DownloadSong(currSong.Url, ytbCookie)
 	if err != nil {
 		return fmt.Errorf("error downloading song: %w", err)
 	}
@@ -144,7 +159,7 @@ func (s *SongList) PlaySong(gid string, cid string, bot *bot.Bot) error {
 
 					// play the next song
 					go func() {
-						err = s.PlaySong(gid, cid, bot)
+						err = s.PlaySong(gid, cid, bot, ytbCookie)
 						if err != nil {
 							fmt.Println("Error playing song:", err)
 						}

internal/music/music_test.go

@@ -1,20 +1,23 @@
 package music
 
 import (
+	"github.com/joho/godotenv"
 	"os"
 	"testing"
 )
 
 func TestDownloadSong(t *testing.T) {
+	err := godotenv.Load("../../.env")
+
 	if err := os.Remove("songCache\\d3J3uJpCgos.mp3"); err != nil {
 		if !os.IsNotExist(err) {
 			t.Errorf("Failed to remove test file: %v", err)
 		}
 	}
 
-	song, err := DownloadSong("https://www.youtube.com/watch?v=d3J3uJpCgos&list=PLwCTYY94JxbZHrJ-anoUuFkNHSFQqe438&index=6&pp=gAQBiAQB8AUB")
+	song, err := DownloadSong("https://www.youtube.com/watch?v=d3J3uJpCgos&list=PLwCTYY94JxbZHrJ-anoUuFkNHSFQqe438&index=6&pp=gAQBiAQB8AUB", os.Getenv("YOUTUBE_COOKIE"))
 	if err != nil {
-		return
+		t.Errorf("Expected no error, got: %v", err)
 	}
 
 	if song == "" {
@@ -27,7 +30,7 @@ func TestDownloadSong(t *testing.T) {
 }
 
 func TestDownloadSongError(t *testing.T) {
-	song, err := DownloadSong("https://www.youtube.com/watch?v=")
+	song, err := DownloadSong("https://www.youtube.com/watch?v=", os.Getenv("YOUTUBE_COOKIE"))
 	if err == nil {
 		t.Errorf("Expected error, got nil")
 	}
@@ -38,7 +41,7 @@ func TestDownloadSongError(t *testing.T) {
 }
 
 func TestDownloadSongAlreadyExists(t *testing.T) {
-	song, err := DownloadSong("https://www.youtube.com/watch?v=d3J3uJpCgos&list=PLwCTYY94JxbZHrJ-anoUuFkNHSFQqe438&index=6&pp=gAQBiAQB8AUB")
+	song, err := DownloadSong("https://www.youtube.com/watch?v=d3J3uJpCgos&list=PLwCTYY94JxbZHrJ-anoUuFkNHSFQqe438&index=6&pp=gAQBiAQB8AUB", os.Getenv("YOUTUBE_COOKIE"))
 	if err != nil {
 		return
 	}

internal/storage/settings.go

@@ -16,6 +16,7 @@ type Settings struct {
 	NewsAPIToken        string `json:"news_api_key"`
 	YoutubeToken        string `json:"youtube_api_key"`
 	ChatHistoryFilePath string `json:"chat_history_file_path"`
+	YoutubeCookies      string `json:"youtube_cookies"`
 }
 
 var Setting Settings

internal/tools/musicfunctions.go

@@ -7,6 +7,7 @@ import (
 	"log"
 	"untitled/internal/bot"
 	"untitled/internal/music"
+	"untitled/internal/storage"
 )
 
 type AddSongArgs struct {
@@ -143,7 +144,7 @@ func skipSong(call openai.ToolCall, s *map[string]map[string]*music.SongList, my
 	currSongList.Songs = currSongList.Songs[1:]
 	currSongList.Mu.Unlock()
 
-	err = currSongList.PlaySong(args.GID, args.CID, myBot)
+	err = currSongList.PlaySong(args.GID, args.CID, myBot, storage.Setting.YoutubeCookies)
 
 	if err != nil {
 		log.Printf("Error playing skipped song: %v", err)
@@ -174,7 +175,7 @@ func playSong(call openai.ToolCall, s *map[string]map[string]*music.SongList, my
 		return `Cannot play song while already playing`, fmt.Errorf("cannot play song while already playing")
 	}
 
-	err = currSongList.PlaySong(args.GID, args.CID, myBot)
+	err = currSongList.PlaySong(args.GID, args.CID, myBot, storage.Setting.YoutubeCookies)
 
 	if err != nil {
 		log.Printf("Error playing song: %v", err)