From da21d6fb1f41c1c5fcc0e6d78da6d2e7ff28544d Mon Sep 17 00:00:00 2001 From: WorldObservationLog Date: Sat, 4 May 2024 22:54:21 +0800 Subject: [PATCH] feat: support ec3 and ac3 codec --- src/cmd.py | 2 +- src/mp4.py | 13 +++++++------ src/rip.py | 2 ++ src/save.py | 2 ++ src/types.py | 6 ++++-- src/utils.py | 4 +++- 6 files changed, 19 insertions(+), 10 deletions(-) diff --git a/src/cmd.py b/src/cmd.py index d633e95..1b5a284 100644 --- a/src/cmd.py +++ b/src/cmd.py @@ -35,7 +35,7 @@ class NewInteractiveShell: download_parser = subparser.add_parser("download") download_parser.add_argument("url", type=str) download_parser.add_argument("-c", "--codec", - choices=["alac", "ec3", "aac", "aac-binaural", "aac-downmix"], default="alac") + choices=["alac", "ec3", "aac", "aac-binaural", "aac-downmix", "ac3"], default="alac") download_parser.add_argument("-f", "--force", type=bool, default=False) subparser.add_parser("exit") diff --git a/src/mp4.py b/src/mp4.py index 74488cc..94d2195 100644 --- a/src/mp4.py +++ b/src/mp4.py @@ -21,8 +21,6 @@ async def extract_media(m3u8_url: str, codec: str) -> Tuple[str, list[str], str] if not specifyPlaylist: raise CodecNotFoundException selected_codec = specifyPlaylist.media[0].group_id - if not specifyPlaylist: - raise stream = m3u8.load(specifyPlaylist.absolute_uri) skds = [key.uri for key in stream.keys if regex.match('(skd?://[^"]*)', key.uri)] keys = [prefetchKey] @@ -30,7 +28,7 @@ async def extract_media(m3u8_url: str, codec: str) -> Tuple[str, list[str], str] match codec: case Codec.ALAC: key_suffix = CodecKeySuffix.KeySuffixAlac - case Codec.EC3: + case Codec.EC3 | Codec.AC3: key_suffix = CodecKeySuffix.KeySuffixAtmos case Codec.AAC: key_suffix = CodecKeySuffix.KeySuffixAAC @@ -103,6 +101,8 @@ def encapsulate(song_info: SongInfo, decrypted_media: bytes, atmos_convent: bool f.write(decrypted_media) if song_info.codec == Codec.EC3 and not atmos_convent: song_name = Path(tmp_dir.name) / Path(name).with_suffix(".ec3") + elif song_info.codec == Codec.AC3 and not atmos_convent: + song_name = Path(tmp_dir.name) / Path(name).with_suffix(".ac3") else: song_name = Path(tmp_dir.name) / Path(name).with_suffix(".m4a") match song_info.codec: @@ -122,12 +122,13 @@ def encapsulate(song_info: SongInfo, decrypted_media: bytes, atmos_convent: bool f"mp4edit --insert moov/trak/mdia/minf/stbl/stsd/alac:{alac_params_atom_name.absolute()} {song_name.absolute()} {final_m4a_name.absolute()}", stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) song_name = final_m4a_name - case Codec.EC3: + case Codec.EC3 | Codec.AC3: if not atmos_convent: with open(song_name.absolute(), "wb") as f: f.write(decrypted_media) - subprocess.run(f"gpac -i {media.absolute()} -o {song_name.absolute()}", - stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + else: + subprocess.run(f"gpac -i {media.absolute()} -o {song_name.absolute()}", + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) case Codec.AAC_BINAURAL | Codec.AAC_DOWNMIX | Codec.AAC: nhml_name = Path(tmp_dir.name) / Path(f"{name}.nhml") with open(nhml_name.absolute(), "w", encoding="utf-8") as f: diff --git a/src/rip.py b/src/rip.py index 7138ca5..5f60fe4 100644 --- a/src/rip.py +++ b/src/rip.py @@ -38,6 +38,8 @@ async def rip_song(song: Song, auth_params: GlobalAuthParams, codec: str, config song = encapsulate(song_info, decrypted_song, config.download.atmosConventToM4a) if codec != Codec.EC3 or (codec == Codec.EC3 and config.download.atmosConventToM4a): song = write_metadata(song, song_metadata, config.metadata.embedMetadata, config.download.coverFormat) + if codec != Codec.AC3 or (codec == Codec.AC3 and config.download.atmosConventToM4a): + song = write_metadata(song, song_metadata, config.metadata.embedMetadata, config.download.coverFormat) save(song, codec, song_metadata, config.download) logger.info(f"Song {song_metadata.artist} - {song_metadata.title} saved!") diff --git a/src/save.py b/src/save.py index 75719b1..119c087 100644 --- a/src/save.py +++ b/src/save.py @@ -14,6 +14,8 @@ def save(song: bytes, codec: str, metadata: SongMetadata, config: Download): os.makedirs(dir_path.absolute()) if codec == Codec.EC3 and not config.atmosConventToM4a: song_path = dir_path / Path(song_name).with_suffix(".ec3") + elif codec == Codec.AC3 and not config.atmosConventToM4a: + song_path = dir_path / Path(song_name).with_suffix(".ac3") else: song_path = dir_path / Path(song_name).with_suffix(".m4a") with open(song_path.absolute(), "wb") as f: diff --git a/src/types.py b/src/types.py index 26a4791..17ef677 100644 --- a/src/types.py +++ b/src/types.py @@ -23,6 +23,7 @@ class SongInfo(BaseModel): class Codec: ALAC = "alac" EC3 = "ec3" + AC3 = "ac3" AAC_BINAURAL = "aac-binaural" AAC_DOWNMIX = "aac-downmix" AAC = "aac" @@ -38,7 +39,8 @@ class CodecKeySuffix: class CodecRegex: - RegexCodecAtmos = "audio-atmos-\\d{4}$" + RegexCodecAtmos = "audio-(atmos|ec3)-\\d{4}$" + RegexCodecAC3 = "audio-ac3-\\d{3}$" RegexCodecAlac = "audio-alac-stereo-\\d{5}-\\d{2}$" RegexCodecBinaural = "audio-stereo-\\d{3}-binaural$" RegexCodecDownmix = "audio-stereo-\\d{3}-downmix$" @@ -48,7 +50,7 @@ class CodecRegex: def get_pattern_by_codec(cls, codec: str): codec_pattern_mapping = {Codec.ALAC: cls.RegexCodecAlac, Codec.EC3: cls.RegexCodecAtmos, Codec.AAC_DOWNMIX: cls.RegexCodecDownmix, Codec.AAC_BINAURAL: cls.RegexCodecBinaural, - Codec.AAC: cls.RegexCodecAAC} + Codec.AAC: cls.RegexCodecAAC, Codec.AC3: cls.RegexCodecAC3} return codec_pattern_mapping.get(codec) diff --git a/src/utils.py b/src/utils.py index 29b4b0e..7d7762c 100644 --- a/src/utils.py +++ b/src/utils.py @@ -109,9 +109,11 @@ def check_song_exists(metadata, config: Download, codec: str): dir_path = Path(config.dirPathFormat.format(**metadata.model_dump())) if not config.atmosConventToM4a and codec == Codec.EC3: return (Path(dir_path) / Path(song_name).with_suffix(".ec3")).exists() + elif not config.atmosConventToM4a and codec == Codec.AC3: + return (Path(dir_path) / Path(song_name).with_suffix(".ac3")).exists() else: return (Path(dir_path) / Path(song_name).with_suffix(".m4a")).exists() def get_valid_filename(filename: str): - return "".join(i for i in filename if i not in "\/:*?<>|") + return "".join(i for i in filename if i not in r"\/:*?<>|")