From 48a8bf3325ad9cf831d067b13f3a08bfcb8aa8f0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B6=B5=E6=9B=A6?= <im.hanxi@gmail.com>
Date: Fri, 6 Dec 2024 22:55:14 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20tag=20=E4=BF=A1=E6=81=AF=E6=94=AF?=
 =?UTF-8?q?=E6=8C=81=E5=86=99=E5=85=A5=E5=88=B0=E6=AD=8C=E6=9B=B2=E6=96=87?=
 =?UTF-8?q?=E4=BB=B6=20see=20#266?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 xiaomusic/config.py                   |   3 +
 xiaomusic/static/default/setting.html |   6 ++
 xiaomusic/utils.py                    | 128 +++++++++++++++++++++++++-
 xiaomusic/xiaomusic.py                |   5 +-
 4 files changed, 140 insertions(+), 2 deletions(-)

diff --git a/xiaomusic/config.py b/xiaomusic/config.py
index 50d1176..933633f 100644
--- a/xiaomusic/config.py
+++ b/xiaomusic/config.py
@@ -168,6 +168,9 @@ class Config:
     enable_yt_dlp_cookies: bool = (
         os.getenv("XIAOMUSIC_ENABLE_YT_DLP_COOKIES", "false").lower() == "true"
     )
+    enable_save_tag: bool = (
+        os.getenv("XIAOMUSIC_ENABLE_SAVE_TAG", "false").lower() == "true"
+    )
     get_ask_by_mina: bool = (
         os.getenv("XIAOMUSIC_GET_ASK_BY_MINA", "false").lower() == "true"
     )
diff --git a/xiaomusic/static/default/setting.html b/xiaomusic/static/default/setting.html
index bf845a6..8b623cc 100644
--- a/xiaomusic/static/default/setting.html
+++ b/xiaomusic/static/default/setting.html
@@ -198,6 +198,12 @@ var vConsole = new window.VConsole();
           <option value="false" selected>false</option>
         </select>
 
+        <label for="enable_save_tag">启用ID3标签写入文件:</label>
+        <select id="enable_save_tag">
+          <option value="true">true</option>
+          <option value="false" selected>false</option>
+        </select>
+
         <label for="get_ask_by_mina">特殊型号获取对话记录:</label>
         <select id="get_ask_by_mina">
           <option value="true">true</option>
diff --git a/xiaomusic/utils.py b/xiaomusic/utils.py
index 2f58ad6..3d969c4 100644
--- a/xiaomusic/utils.py
+++ b/xiaomusic/utils.py
@@ -27,7 +27,19 @@ import aiohttp
 import mutagen
 from mutagen.asf import ASF
 from mutagen.flac import FLAC
-from mutagen.id3 import APIC, ID3, Encoding, TextFrame, TimeStampTextFrame
+from mutagen.id3 import (
+    APIC,
+    ID3,
+    TALB,
+    TCON,
+    TDRC,
+    TIT2,
+    TPE1,
+    USLT,
+    Encoding,
+    TextFrame,
+    TimeStampTextFrame,
+)
 from mutagen.mp3 import MP3
 from mutagen.mp4 import MP4
 from mutagen.oggvorbis import OggVorbis
@@ -568,6 +580,16 @@ class Metadata:
     picture: str = ""
     lyrics: str = ""
 
+    def __init__(self, info=None):
+        if info:
+            self.title = info.get("title", "")
+            self.artist = info.get("artist", "")
+            self.album = info.get("album", "")
+            self.year = info.get("year", "")
+            self.genre = info.get("genre", "")
+            self.picture = info.get("picture", "")
+            self.lyrics = info.get("lyrics", "")
+
 
 def _get_alltag_value(tags, k):
     v = tags.getall(k)
@@ -731,6 +753,110 @@ def extract_audio_metadata(file_path, save_root):
     return asdict(metadata)
 
 
+def set_music_tag_to_file(file_path, info):
+    audio = mutagen.File(file_path, easy=True)
+    if audio is None:
+        log.error(f"Unable to open file {file_path}")
+        return "Unable to open file"
+
+    if isinstance(audio, MP3):
+        _set_mp3_tags(audio, info)
+    elif isinstance(audio, FLAC):
+        _set_flac_tags(audio, info)
+    elif isinstance(audio, MP4):
+        _set_mp4_tags(audio, info)
+    elif isinstance(audio, OggVorbis):
+        _set_ogg_tags(audio, info)
+    elif isinstance(audio, ASF):
+        _set_asf_tags(audio, info)
+    elif isinstance(audio, WAVE):
+        _set_wave_tags(audio, info)
+    else:
+        log.error(f"Unsupported file type for {file_path}")
+        return "Unsupported file type"
+
+    try:
+        audio.save()
+        log.info(f"Tags saved successfully to {file_path}")
+        return "OK"
+    except Exception as e:
+        log.exception(f"Error saving tags: {e}")
+        return "Error saving tags"
+
+
+def _set_mp3_tags(audio, info):
+    audio["TIT2"] = TIT2(encoding=3, text=info.title)
+    audio["TPE1"] = TPE1(encoding=3, text=info.artist)
+    audio["TALB"] = TALB(encoding=3, text=info.album)
+    audio["TDRC"] = TDRC(encoding=3, text=info.year)
+    audio["TCON"] = TCON(encoding=3, text=info.genre)
+    if info.lyrics:
+        audio["USLT"] = USLT(encoding=3, lang="eng", text=info.lyrics)
+    if info.picture:
+        with open(info.picture, "rb") as img_file:
+            image_data = img_file.read()
+        audio["APIC"] = APIC(
+            encoding=3, mime="image/jpeg", type=3, desc="Cover", data=image_data
+        )
+
+
+def _set_flac_tags(audio, info):
+    audio["TITLE"] = info.title
+    audio["ARTIST"] = info.artist
+    audio["ALBUM"] = info.album
+    audio["DATE"] = info.year
+    audio["GENRE"] = info.genre
+    if info.lyrics:
+        audio["LYRICS"] = info.lyrics
+    if info.picture:
+        with open(info.picture, "rb") as img_file:
+            image_data = img_file.read()
+        audio.add_picture(image_data)
+
+
+def _set_mp4_tags(audio, info):
+    audio["\xa9nam"] = info.title
+    audio["\xa9ART"] = info.artist
+    audio["\xa9alb"] = info.album
+    audio["\xa9day"] = info.year
+    audio["\xa9gen"] = info.genre
+    if info.picture:
+        with open(info.picture, "rb") as img_file:
+            image_data = img_file.read()
+        audio["covr"] = [image_data]
+
+
+def _set_ogg_tags(audio, info):
+    audio["TITLE"] = info.title
+    audio["ARTIST"] = info.artist
+    audio["ALBUM"] = info.album
+    audio["DATE"] = info.year
+    audio["GENRE"] = info.genre
+    if info.lyrics:
+        audio["LYRICS"] = info.lyrics
+    if info.picture:
+        with open(info.picture, "rb") as img_file:
+            image_data = img_file.read()
+        audio["metadata_block_picture"] = base64.b64encode(image_data).decode()
+
+
+def _set_asf_tags(audio, info):
+    audio["Title"] = info.title
+    audio["Author"] = info.artist
+    audio["WM/AlbumTitle"] = info.album
+    audio["WM/Year"] = info.year
+    audio["WM/Genre"] = info.genre
+    if info.picture:
+        with open(info.picture, "rb") as img_file:
+            image_data = img_file.read()
+        audio["WM/Picture"] = image_data
+
+
+def _set_wave_tags(audio, info):
+    audio["Title"] = info.title
+    audio["Artist"] = info.artist
+
+
 # 下载播放列表
 async def download_playlist(config, url, dirname):
     title = f"{dirname}/%(title)s.%(ext)s"
diff --git a/xiaomusic/xiaomusic.py b/xiaomusic/xiaomusic.py
index a87c114..5ab1ed9 100644
--- a/xiaomusic/xiaomusic.py
+++ b/xiaomusic/xiaomusic.py
@@ -52,6 +52,7 @@ from xiaomusic.utils import (
     parse_cookie_string,
     parse_str_to_dict,
     save_picture_by_base64,
+    set_music_tag_to_file,
     traverse_music_directory,
     try_add_access_control_param,
 )
@@ -477,11 +478,13 @@ class XiaoMusic:
         tags["year"] = info.year
         tags["genre"] = info.genre
         tags["lyrics"] = info.lyrics
+        file_path = self.all_music[name]
         if info.picture:
-            file_path = self.all_music[name]
             tags["picture"] = save_picture_by_base64(
                 info.picture, self.config.picture_cache_path, file_path
             )
+        if self.config.enable_save_tag and (not self.is_web_music(name)):
+            set_music_tag_to_file(file_path, Metadata(tags))
         self.all_music_tags[name] = tags
         self.try_save_tag_cache()
         return "OK"