diff --git a/config.go b/config.go index b1de0b3..ab97ea7 100644 --- a/config.go +++ b/config.go @@ -17,6 +17,7 @@ type Config struct { DownloadLocation string Token string TempDirLocation string + NumWorkers int Slides bool } @@ -37,6 +38,14 @@ func parseConfig(configLocation string) *Config { panic(err) } + if config.TempDirLocation == "" { + config.TempDirLocation = "./temp" + } + + if config.NumWorkers == 0 { + config.NumWorkers = 5 + } + return &config } diff --git a/impartus.go b/impartus.go index d0b0fd2..85c7fb9 100644 --- a/impartus.go +++ b/impartus.go @@ -311,10 +311,10 @@ func downloadUrl(url string, id int, chunk int, view string) (string, error) { outFilepath := filepath.Join(config.TempDirLocation, fmt.Sprintf("%d_%s_%04d.ts.temp", id, view, chunk)) outFile, err := os.Create(outFilepath) - defer outFile.Close() if err != nil { fmt.Printf("Could not download chunk %d %v", chunk, err) } + defer outFile.Close() _, err = io.Copy(outFile, resp.Body) if err != nil { @@ -331,62 +331,52 @@ type DownloadedPlaylist struct { Playlist ParsedPlaylist } -func DownloadPlaylist(playlists []ParsedPlaylist) []DownloadedPlaylist { +func DownloadPlaylist(playlist ParsedPlaylist) DownloadedPlaylist { config := GetConfig() - var downloaded []DownloadedPlaylist + var downloadedPlaylist DownloadedPlaylist - err := os.MkdirAll(config.TempDirLocation, 0755) + resp, _ := GetClientAuthorized(playlist.KeyURL, config.Token) + defer resp.Body.Close() + keyUrlContent, err := io.ReadAll(resp.Body) if err != nil { - log.Fatalln("Could not create temp dir") + fmt.Printf("Could not get keyUrlContent %v", err) + panic(err) } - for _, playlist := range playlists { - var downloadedPlaylist DownloadedPlaylist - - resp, _ := GetClientAuthorized(playlist.KeyURL, config.Token) - defer resp.Body.Close() - keyUrlContent, err := io.ReadAll(resp.Body) - if err != nil { - fmt.Printf("Could not get keyUrlContent %v", err) - panic(err) - } + decryptionKey := getDecryptionKey(keyUrlContent) - decryptionKey := getDecryptionKey(keyUrlContent) - - if len(playlist.FirstViewURLs) > 0 && config.Views != "left" { - bar := progressbar.NewOptions64(-1, progressbar.OptionSetDescription(fmt.Sprintf("Lec %03d downloading right view chunks", playlist.SeqNo))) - for i, url := range playlist.FirstViewURLs { - f, err := downloadUrl(url, playlist.Id, i, "first") - if err != nil { - fmt.Println() - fmt.Println("[WARNING] Chunk", i, "download failed") - continue - } - chunkPath := decryptChunk(f, decryptionKey) - downloadedPlaylist.FirstViewChunks = append(downloadedPlaylist.FirstViewChunks, chunkPath) - bar.Add(1) + if len(playlist.FirstViewURLs) > 0 && config.Views != "left" { + bar := progressbar.NewOptions64(-1, progressbar.OptionSetDescription(fmt.Sprintf("Lec %03d downloading right view chunks", playlist.SeqNo))) + for i, url := range playlist.FirstViewURLs { + f, err := downloadUrl(url, playlist.Id, i, "first") + if err != nil { + fmt.Println() + fmt.Println("[WARNING] Chunk", i, "download failed") + continue } + chunkPath := decryptChunk(f, decryptionKey) + downloadedPlaylist.FirstViewChunks = append(downloadedPlaylist.FirstViewChunks, chunkPath) + bar.Add(1) } + } - if len(playlist.SecondViewURLs) > 0 && config.Views != "right" { - bar := progressbar.NewOptions64(-1, progressbar.OptionSetDescription(fmt.Sprintf("Lec %03d downloading left view chunks", playlist.SeqNo))) - for i, url := range playlist.SecondViewURLs { - f, err := downloadUrl(url, playlist.Id, i, "second") - if err != nil { - fmt.Println() - fmt.Println("[WARNING] Chunk", i, "download failed") - continue - } - chunkPath := decryptChunk(f, decryptionKey) - downloadedPlaylist.SecondViewChunks = append(downloadedPlaylist.SecondViewChunks, chunkPath) - bar.Add(1) + if len(playlist.SecondViewURLs) > 0 && config.Views != "right" { + bar := progressbar.NewOptions64(-1, progressbar.OptionSetDescription(fmt.Sprintf("Lec %03d downloading left view chunks", playlist.SeqNo))) + for i, url := range playlist.SecondViewURLs { + f, err := downloadUrl(url, playlist.Id, i, "second") + if err != nil { + fmt.Println() + fmt.Println("[WARNING] Chunk", i, "download failed") + continue } + chunkPath := decryptChunk(f, decryptionKey) + downloadedPlaylist.SecondViewChunks = append(downloadedPlaylist.SecondViewChunks, chunkPath) + bar.Add(1) } - - downloadedPlaylist.Playlist = playlist - downloaded = append(downloaded, downloadedPlaylist) } - return downloaded + + downloadedPlaylist.Playlist = playlist + return downloadedPlaylist } type M3U8File struct { @@ -395,66 +385,61 @@ type M3U8File struct { Playlist ParsedPlaylist } -func CreateTempM3U8Files(downloadedPlaylists []DownloadedPlaylist) []M3U8File { - var m3u8Files []M3U8File - +func CreateTempM3U8File(downloadedPlaylist DownloadedPlaylist) M3U8File { config := GetConfig() - for _, playlist := range downloadedPlaylists { - var m3u8File M3U8File - - if len(playlist.FirstViewChunks) > 0 { - firstView, err := os.Create(fmt.Sprintf("%s/%d_first.m3u8", config.TempDirLocation, playlist.Playlist.Id)) - defer firstView.Close() + var m3u8File M3U8File - if err != nil { - fmt.Printf("Could not create temp m3u8 file for ttid %d with error %v", playlist.Playlist.Id, err) - } + if len(downloadedPlaylist.FirstViewChunks) > 0 { + firstView, err := os.Create(fmt.Sprintf("%s/%d_first.m3u8", config.TempDirLocation, downloadedPlaylist.Playlist.Id)) + if err != nil { + fmt.Printf("Could not create temp m3u8 file for ttid %d with error %v", downloadedPlaylist.Playlist.Id, err) + } + defer firstView.Close() - firstView.WriteString(`#EXTM3U + firstView.WriteString(`#EXTM3U #EXT-X-VERSION:3 #EXT-X-MEDIA-SEQUENCE:0 #EXT-X-ALLOW-CACHE:YES #EXT-X-TARGETDURATION:11 #EXT-X-KEY:METHOD=NONE`) - for _, chunk := range playlist.FirstViewChunks { - firstView.WriteString("#EXTINF:1\n") - firstView.WriteString("../" + chunk + "\n") - } + for _, chunk := range downloadedPlaylist.FirstViewChunks { + firstView.WriteString("#EXTINF:1\n") + firstView.WriteString("../" + chunk + "\n") + } - firstView.WriteString("#EXT-X-ENDLIST") + firstView.WriteString("#EXT-X-ENDLIST") - m3u8File.FirstViewFile = firstView.Name() - } + m3u8File.FirstViewFile = firstView.Name() + } - if len(playlist.SecondViewChunks) > 0 { - secondView, err := os.Create(fmt.Sprintf("%s/%d_second.m3u8", config.TempDirLocation, playlist.Playlist.Id)) - if err != nil { - fmt.Printf("Could not create temp m3u8 file for ttid %d with error %v", playlist.Playlist.Id, err) - } - defer secondView.Close() + if len(downloadedPlaylist.SecondViewChunks) > 0 { + secondView, err := os.Create(fmt.Sprintf("%s/%d_second.m3u8", config.TempDirLocation, downloadedPlaylist.Playlist.Id)) + if err != nil { + fmt.Printf("Could not create temp m3u8 file for ttid %d with error %v", downloadedPlaylist.Playlist.Id, err) + } + defer secondView.Close() - secondView.WriteString(`#EXTM3U + secondView.WriteString(`#EXTM3U #EXT-X-VERSION:3 #EXT-X-MEDIA-SEQUENCE:0 #EXT-X-ALLOW-CACHE:YES #EXT-X-TARGETDURATION:11 #EXT-X-KEY:METHOD=NONE`) - for _, chunk := range playlist.SecondViewChunks { - secondView.WriteString("#EXTINF:1\n") - secondView.WriteString("../" + chunk + "\n") - } - - secondView.WriteString("#EXT-X-ENDLIST") - - m3u8File.SecondViewFile = secondView.Name() + for _, chunk := range downloadedPlaylist.SecondViewChunks { + secondView.WriteString("#EXTINF:1\n") + secondView.WriteString("../" + chunk + "\n") } - m3u8File.Playlist = playlist.Playlist - m3u8Files = append(m3u8Files, m3u8File) + secondView.WriteString("#EXT-X-ENDLIST") + + m3u8File.SecondViewFile = secondView.Name() } - return m3u8Files + + m3u8File.Playlist = downloadedPlaylist.Playlist + + return m3u8File } func DownloadLectureSlides(lecture Lecture) { diff --git a/main.go b/main.go index eacbb43..e08a3ca 100644 --- a/main.go +++ b/main.go @@ -7,6 +7,7 @@ import ( "os/exec" "path/filepath" "strings" + "sync" ) func main() { @@ -58,25 +59,54 @@ func main() { } } - downloadedPlaylists := DownloadPlaylist(GetPlaylist(chosenLectures)) - metadataFiles := CreateTempM3U8Files(downloadedPlaylists) + playlists := GetPlaylist(chosenLectures) - for _, file := range metadataFiles { - var left, right string - if file.FirstViewFile != "" && config.Views != "left" { - left = JoinChunksFromM3U8(file.FirstViewFile, fmt.Sprintf("LEC %03d %s RIGHT VIEW.mp4", startLectureIndex+1, file.Playlist.Title)) - } + err = os.MkdirAll(config.TempDirLocation, 0755) + if err != nil { + log.Fatalln("Could not create temp dir") + } - if file.SecondViewFile != "" && config.Views != "right" { - right = JoinChunksFromM3U8(file.SecondViewFile, fmt.Sprintf("LEC %03d %s LEFT VIEW.mp4", startLectureIndex+1, file.Playlist.Title)) - } + numWorkers := config.NumWorkers + playlistJobs := make(chan ParsedPlaylist, numWorkers) - if left != "" && right != "" && config.Views == "both" { - JoinViews(left, right, fmt.Sprintf("LEC %03d %s", startLectureIndex+1, file.Playlist.Title)) - } - startLectureIndex++ + var joinWg sync.WaitGroup + for i := 0; i < numWorkers; i++ { + go func() { + for playlist := range playlistJobs { + fmt.Println("Downloading playlist: ", playlist.Title, playlist.SeqNo) + downloadedPlaylist := DownloadPlaylist(playlist) + metadataFile := CreateTempM3U8File(downloadedPlaylist) + fmt.Println("Downloaded playlist: ", playlist.Title, playlist.SeqNo) + + go func(file M3U8File) { + defer joinWg.Done() + fmt.Println("Joining chunks for: ", file.Playlist.Title, file.Playlist.SeqNo) + var left, right string + if file.FirstViewFile != "" && config.Views != "left" { + left = JoinChunksFromM3U8(file.FirstViewFile, fmt.Sprintf("LEC %03d %s RIGHT VIEW.mp4", file.Playlist.SeqNo, file.Playlist.Title)) + } + if file.SecondViewFile != "" && config.Views != "right" { + right = JoinChunksFromM3U8(file.SecondViewFile, fmt.Sprintf("LEC %03d %s LEFT VIEW.mp4", file.Playlist.SeqNo, file.Playlist.Title)) + } + + if left != "" && right != "" && config.Views == "both" { + JoinViews(left, right, fmt.Sprintf("LEC %03d %s", file.Playlist.SeqNo, file.Playlist.Title)) + } + fmt.Println("Joined chunks for: ", file.Playlist.Title, file.Playlist.SeqNo) + }(metadataFile) + } + }() } + for _, playlist := range playlists { + fmt.Println("Adding playlist to job queue: ", playlist.Title, playlist.SeqNo) + joinWg.Add(1) + playlistJobs <- playlist + } + + joinWg.Wait() + close(playlistJobs) + fmt.Print("\n\n") fmt.Println("It is recommended that you use this tool as sparingly as possible. Heavy usage of this tool puts more strain on impartus server leading to potential IP bans, breakingAPI changes and possibly legal action.") fmt.Println("If this project helped you, consider starring it on GitHub: https://github.com/pnicto/impartus-video-downloader") diff --git a/parser.go b/parser.go index 3f47e72..a78f88a 100644 --- a/parser.go +++ b/parser.go @@ -8,18 +8,18 @@ import ( ) type ParsedPlaylist struct { - HasMultipleViews bool KeyURL string + Title string FirstViewURLs []string SecondViewURLs []string Id int - Title string SeqNo int + HasMultipleViews bool } func PlaylistParser(scanner *bufio.Scanner, id int, title string, seqNo int) ParsedPlaylist { var parsedOutput ParsedPlaylist - var isFirstView = true + isFirstView := true var firstViewUrls []string var secondViewUrls []string diff --git a/sample.config.json b/sample.config.json index d1cfa9c..a80a4ba 100644 --- a/sample.config.json +++ b/sample.config.json @@ -6,5 +6,6 @@ "views": "both", "downloadLocation": "./downloads", "tempDirLocation": "./.temp", - "slides": false + "slides": false, + "numWorkers": 5 }