add atmos and select to main.go

--atmos,--select

main.go

@@ -37,6 +37,8 @@ const (
 var (
 	forbiddenNames = regexp.MustCompile(`[/\\<>:"|?*]`)
 )
+var dl_atmos = false
+var dl_select = false
 
 type Config struct {
 	MediaUserToken          string `yaml:"media-user-token"`
@@ -63,6 +65,7 @@ type Config struct {
 	GetM3u8Port             string `yaml:"get-m3u8-port"`
 	GetM3u8FromDevice       bool   `yaml:"get-m3u8-from-device"`
 	AlacMax                 int    `yaml:"alac-max"`
+	AtmosMax                int    `yaml:"atmos-max"`
 	UseSongInfoForPlaylist  bool   `yaml:"use-songinfo-for-playlist"`
 	DlAlbumcoverForPlaylist bool   `yaml:"dl-albumcover-for-playlist"`
 }
@@ -109,6 +112,15 @@ func (*Alac) GetType() mp4.BoxType {
 	return BoxTypeAlac()
 }
 
+func isInArray(arr []int, target int) bool {
+	for _, num := range arr {
+		if num == target {
+			return true
+		}
+	}
+	return false
+}
+
 func fileExists(path string) (bool, error) {
 	f, err := os.Stat(path)
 	if err == nil {
@@ -1009,6 +1021,14 @@ func decryptSong(info *SongInfo, keys []string, manifest *AutoGenerated, filenam
 		return err
 	}
 	defer create.Close()
+	if dl_atmos {
+		_, err = create.Write(decrypted)
+		if err != nil {
+			panic(err)
+		}
+
+		return nil
+	}
 
 	return writeM4a(mp4.NewWriter(create), info, manifest, decrypted, trackNum, trackTotal)
 }
@@ -1304,6 +1324,12 @@ func writeLyrics(sanAlbumFolder, filename string, lrc string) error {
 }
 
 func rip(albumId string, token string, storefront string, userToken string) error {
+	var Codec string
+	if dl_atmos {
+		Codec = "Atmos"
+	} else {
+		Codec = "ALAC"
+	}
 	meta, err := getMeta(albumId, token, storefront)
 	if err != nil {
 		fmt.Println("Failed to get album metadata.\n")
@@ -1336,12 +1362,15 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 	singerFolder := filepath.Join(config.AlacSaveFolder, forbiddenNames.ReplaceAllString(singerFoldername, "_"))
 	var Quality string
 	if strings.Contains(config.AlbumFolderFormat, "Quality") {
+		if dl_atmos {
+			Quality = fmt.Sprintf("%dkbps", config.AtmosMax-2000)
+		} else {
 			manifest1, err := getInfoFromAdam(meta.Data[0].Relationships.Tracks.Data[0].ID, token, storefront)
 			if err != nil {
 				fmt.Println("Failed to get manifest.\n", err)
 			} else {
 				if manifest1.Attributes.ExtendedAssetUrls.EnhancedHls == "" {
-				fmt.Println("Unavailable in ALAC.\n")
+					fmt.Println("Unavailable.\n")
 				} else {
 					EnhancedHls_m3u8, err := checkM3u8(meta.Data[0].Relationships.Tracks.Data[0].ID, "album")
 					if strings.HasPrefix(EnhancedHls_m3u8, "http") {
@@ -1354,6 +1383,7 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 				}
 			}
 		}
+	}
 	stringsToJoin := []string{}
 	if meta.Data[0].Attributes.IsAppleDigitalMaster || meta.Data[0].Attributes.IsMasteredForItunes {
 		if config.AppleMasterChoice != "" {
@@ -1378,7 +1408,7 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 			"{PlaylistName}", meta.Data[0].Attributes.Name,
 			"{PlaylistId}", albumId,
 			"{Quality}", Quality,
-			"{Codec}", "ALAC",
+			"{Codec}", Codec,
 			"{Tag}", Tag_string,
 		).Replace(config.PlaylistFolderFormat)
 	} else {
@@ -1392,7 +1422,7 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 			"{Copyright}", meta.Data[0].Attributes.Copyright,
 			"{AlbumId}", albumId,
 			"{Quality}", Quality,
-			"{Codec}", "ALAC",
+			"{Codec}", Codec,
 			"{Tag}", Tag_string,
 		).Replace(config.AlbumFolderFormat)
 	}
@@ -1448,8 +1478,51 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 		}
 	}
 	trackTotal := len(meta.Data[0].Relationships.Tracks.Data)
+	arr := make([]int, trackTotal)
+	for i := 0; i < trackTotal; i++ {
+		arr[i] = i + 1
+	}
+	selected := []int{}
+
+	if !dl_select {
+		selected = arr
+	} else {
+
+		fmt.Print("select: ")
+		reader := bufio.NewReader(os.Stdin)
+		input, err := reader.ReadString('\n')
+		if err != nil {
+			fmt.Println(err)
+		}
+		input = strings.TrimSpace(input)
+		inputs := strings.Fields(input)
+
+		for _, str := range inputs {
+			num, err := strconv.Atoi(str)
+			if err != nil {
+				fmt.Printf("wrong '%s', skip...\n", str)
+				continue
+			}
+
+			found := false
+			for i := 0; i < len(arr); i++ {
+				if arr[i] == num {
+					selected = append(selected, num)
+					found = true
+					break
+				}
+			}
+
+			if !found {
+				fmt.Printf("Option '%d' not found or already selected, skipping...\n", num)
+			}
+		}
+
+		fmt.Println("Selected options:", selected)
+	}
 	for trackNum, track := range meta.Data[0].Relationships.Tracks.Data {
 		trackNum++
+		if isInArray(selected, trackNum) {
 			trackTotalnum += 1
 			fmt.Printf("Track %d of %d:\n", trackNum, trackTotal)
 			manifest, err := getInfoFromAdam(track.ID, token, storefront)
@@ -1458,7 +1531,7 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 				continue
 			}
 			if manifest.Attributes.ExtendedAssetUrls.EnhancedHls == "" {
-			fmt.Println("Unavailable in ALAC.")
+				fmt.Println("Unavailable.")
 				continue
 			}
 			EnhancedHls_m3u8, err := checkM3u8(track.ID, "song")
@@ -1467,12 +1540,16 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 			}
 			var Quality string
 			if strings.Contains(config.SongFileFormat, "Quality") {
+				if dl_atmos {
+					Quality = fmt.Sprintf("%dkbps", config.AtmosMax-2000)
+				} else {
 					Quality, err = extractMediaQuality(manifest.Attributes.ExtendedAssetUrls.EnhancedHls)
 					if err != nil {
 						fmt.Println("Failed to extract quality from manifest.\n", err)
 						continue
 					}
 				}
+			}
 			stringsToJoin := []string{}
 			if track.Attributes.IsAppleDigitalMaster {
 				if config.AppleMasterChoice != "" {
@@ -1499,12 +1576,17 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 				"{TrackNumber}", fmt.Sprintf("%0d", track.Attributes.TrackNumber),
 				"{Quality}", Quality,
 				"{Tag}", Tag_string,
-			"{Codec}", "ALAC",
+				"{Codec}", Codec,
 			).Replace(config.SongFileFormat)
 			fmt.Println(songName)
 			filename := fmt.Sprintf("%s.m4a", forbiddenNames.ReplaceAllString(songName, "_"))
+			if dl_atmos {
+				filename = fmt.Sprintf("%s.ec3", forbiddenNames.ReplaceAllString(songName, "_"))
+			}
+			m4afilename := fmt.Sprintf("%s.m4a", forbiddenNames.ReplaceAllString(songName, "_"))
 			lrcFilename := fmt.Sprintf("%s.lrc", forbiddenNames.ReplaceAllString(songName, "_"))
 			trackPath := filepath.Join(sanAlbumFolder, filename)
+			m4atrackPath := filepath.Join(sanAlbumFolder, m4afilename)
 			var lrc string = ""
 			if userToken != "your-media-user-token" && (config.EmbedLrc || config.SaveLrcFile) {
 				ttml, err := getSongLyrics(track.ID, storefront, token, userToken)
@@ -1536,6 +1618,15 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 				oktrackNum += 1
 				continue
 			}
+			m4aexists, err := fileExists(m4atrackPath)
+			if err != nil {
+				fmt.Println("Failed to check if track exists.")
+			}
+			if m4aexists {
+				fmt.Println("Track already exists locally.")
+				oktrackNum += 1
+				continue
+			}
 
 			trackUrl, keys, err := extractMedia(manifest.Attributes.ExtendedAssetUrls.EnhancedHls)
 			if err != nil {
@@ -1568,6 +1659,36 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 			tags := []string{
 				fmt.Sprintf("lyrics=%s", lrc),
 			}
+
+			index := trackNum - 1
+			if dl_atmos {
+				tags = []string{
+					"tool=",
+					fmt.Sprintf("lyrics=%s", lrc),
+					fmt.Sprintf("title=%s", meta.Data[0].Relationships.Tracks.Data[index].Attributes.Name),
+					fmt.Sprintf("artist=%s", meta.Data[0].Relationships.Tracks.Data[index].Attributes.ArtistName),
+					fmt.Sprintf("genre=%s", meta.Data[0].Relationships.Tracks.Data[index].Attributes.GenreNames[0]),
+					fmt.Sprintf("created=%s", meta.Data[0].Attributes.ReleaseDate),
+					fmt.Sprintf("album_artist=%s", meta.Data[0].Attributes.ArtistName),
+					fmt.Sprintf("composer=%s", meta.Data[0].Relationships.Tracks.Data[index].Attributes.ComposerName),
+					fmt.Sprintf("writer=%s", meta.Data[0].Relationships.Tracks.Data[index].Attributes.ComposerName),
+					fmt.Sprintf("performer=%s", meta.Data[0].Attributes.ArtistName),
+					fmt.Sprintf("copyright=%s", meta.Data[0].Attributes.Copyright),
+					fmt.Sprintf("ISRC=%s", meta.Data[0].Relationships.Tracks.Data[index].Attributes.Isrc),
+					fmt.Sprintf("UPC=%s", meta.Data[0].Attributes.Upc),
+				}
+				if strings.Contains(albumId, "pl.") && !config.UseSongInfoForPlaylist {
+					tags = append(tags, "disk=1/1")
+					tags = append(tags, fmt.Sprintf("track=%s", trackNum))
+					tags = append(tags, fmt.Sprintf("tracknum=%d/%d", trackNum, trackTotal))
+					tags = append(tags, fmt.Sprintf("album=%s", meta.Data[0].Attributes.Name))
+				} else {
+					tags = append(tags, fmt.Sprintf("disk=%d/%d", meta.Data[0].Relationships.Tracks.Data[index].Attributes.DiscNumber, meta.Data[0].Relationships.Tracks.Data[trackTotal-1].Attributes.DiscNumber))
+					tags = append(tags, fmt.Sprintf("track=%s", meta.Data[0].Relationships.Tracks.Data[index].Attributes.TrackNumber))
+					tags = append(tags, fmt.Sprintf("tracknum=%d/%d", meta.Data[0].Relationships.Tracks.Data[index].Attributes.TrackNumber, trackTotal))
+					tags = append(tags, fmt.Sprintf("album=%s", meta.Data[0].Relationships.Tracks.Data[index].Attributes.AlbumName))
+				}
+			}
 			if track.Attributes.ContentRating == "explicit" {
 				tags = append(tags, "rating=1")
 			} else if track.Attributes.ContentRating == "clean" {
@@ -1588,6 +1709,9 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 			}
 			tagsString := strings.Join(tags, ":")
 			cmd := exec.Command("MP4Box", "-itags", tagsString, trackPath)
+			if dl_atmos {
+				cmd = exec.Command("MP4Box", "-add", trackPath, "-name", fmt.Sprintf("1=%s", meta.Data[0].Relationships.Tracks.Data[index].Attributes.Name), "-itags", tagsString, "-brand", "mp42", "-ab", "dby1", m4atrackPath)
+			}
 			if err := cmd.Run(); err != nil {
 				fmt.Printf("Embed failed: %v\n", err)
 				continue
@@ -1598,8 +1722,17 @@ func rip(albumId string, token string, storefront string, userToken string) erro
 					continue
 				}
 			}
+			if dl_atmos {
+				fmt.Printf("Deleting original EC3 file: %s\n", filepath.Base(trackPath))
+				if err := os.Remove(trackPath); err != nil {
+					fmt.Printf("Error deleting file: %v\n", err)
+					continue
+				}
+				fmt.Printf("Successfully processed and deleted %s\n", filepath.Base(trackPath))
+			}
 			oktrackNum += 1
 		}
+	}
 	return err
 }
 
@@ -1614,6 +1747,17 @@ func main() {
 		fmt.Println("Failed to get token.")
 		return
 	}
+	var dlArgs []string
+	for _, arg := range os.Args {
+		if strings.Contains(arg, "--atmos") {
+			dl_atmos = true
+		} else if strings.Contains(arg, "--select") {
+			dl_select = true
+		} else {
+			dlArgs = append(dlArgs, arg)
+		}
+	}
+	os.Args = dlArgs
 	if strings.Contains(os.Args[1], "/artist/") {
 		newArgs, err := checkArtist(os.Args[1], token)
 		if err != nil {
@@ -1882,6 +2026,25 @@ func extractMedia(b string) (string, []string, error) {
 		return master.Variants[i].AverageBandwidth > master.Variants[j].AverageBandwidth
 	})
 	for _, variant := range master.Variants {
+		if dl_atmos {
+			if variant.Codecs == "ec-3" {
+				split := strings.Split(variant.Audio, "-")
+				length := len(split)
+				length_int, err := strconv.Atoi(split[length-1])
+				if err != nil {
+					return "", nil, err
+				}
+				if length_int <= config.AtmosMax {
+					fmt.Printf("%s\n", variant.Audio)
+					streamUrlTemp, err := masterUrl.Parse(variant.URI)
+					if err != nil {
+						panic(err)
+					}
+					streamUrl = streamUrlTemp
+					break
+				}
+			}
+		} else {
 			if variant.Codecs == "alac" {
 				split := strings.Split(variant.Audio, "-")
 				length := len(split)
@@ -1900,8 +2063,9 @@ func extractMedia(b string) (string, []string, error) {
 				}
 			}
 		}
+	}
 	if streamUrl == nil {
-		return "", nil, errors.New("no alac codec found")
+		return "", nil, errors.New("no codec found")
 	}
 	var keys []string
 	keys = append(keys, prefetchKey)
@@ -1909,10 +2073,16 @@ func extractMedia(b string) (string, []string, error) {
 	regex := regexp.MustCompile(`"(skd?://[^"]*)"`)
 	matches := regex.FindAllStringSubmatch(masterString, -1)
 	for _, match := range matches {
+		if dl_atmos {
+			if strings.HasSuffix(match[1], "c24") || strings.HasSuffix(match[1], "c6") {
+				keys = append(keys, match[1])
+			}
+		} else {
 			if strings.HasSuffix(match[1], "c23") || strings.HasSuffix(match[1], "c6") {
 				keys = append(keys, match[1])
 			}
 		}
+	}
 	return streamUrl.String(), keys, nil
 }
 func extractVideo(c string) (string, error) {
@@ -1992,6 +2162,8 @@ func extractSong(url string) (*SongInfo, error) {
 		return nil, err
 	}
 
+	var extracted *SongInfo
+	if !dl_atmos {
 		enca, err := mp4.ExtractBoxWithPayload(f, stbl[0], []mp4.BoxType{
 			mp4.BoxTypeStsd(),
 			mp4.BoxTypeEnca(),
@@ -2005,11 +2177,16 @@ func extractSong(url string) (*SongInfo, error) {
 		if err != nil || len(aalac) != 1 {
 			return nil, err
 		}
-
+		extracted = &SongInfo{
-	extracted := &SongInfo{
 			r:         f,
 			alacParam: aalac[0].Payload.(*Alac),
 		}
+	} else {
+		extracted = &SongInfo{
+			r: f,
+			// alacParam: aalac[0].Payload.(*Alac),
+		}
+	}
 
 	moofs, err := mp4.ExtractBox(f, nil, []mp4.BoxType{
 		mp4.BoxTypeMoof(),