from mutagen.easyid3 import ID3
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad,unpad
import io,sys,pathlib
import re,base64,magic
import mutagen



class XMInfo:
    '''
    const {
        title: s,
        artist: l,
        subtitle: c,
        length: d,
        comment: {
            language: u,
            text: p
        },
        album: h,
        trackNumber: b,
        size: g,
        encodingTechnology: v,
        ISRC: _,
        fileType: y,
        encodedBy: w,
        publisher: k,
        composer: x,
        mediaType: S
    }
    '''
    def __init__(self):
        self.title = ""
        self.artist = ""
        self.album = ""
        self.tracknumber = 0
        self.size = 0
        self.header_size = 0
        self.ISRC = ""
        self.encodedby = ""
        self.encoding_technology = ""

    def iv(self):
        if (self.ISRC != ""):
            return bytes.fromhex(self.ISRC)
        return bytes.fromhex(self.encodedby)

def get_str(x):
    if x is None:
        return ""
    return x

def read_file(x):
    with open(x,"rb") as f:
        return f.read()

# return number of id3 bytes
def get_xm_info(data:bytes):
    # print(EasyID3(io.BytesIO(data)))
    id3 = ID3(io.BytesIO(data),v2_version=3)
    id3value = XMInfo()
    id3value.title = str(id3["TIT2"])
    id3value.album = str(id3["TALB"])
    id3value.artist = str(id3["TPE1"])
    id3value.tracknumber = int(str(id3["TRCK"]))
    id3value.ISRC = "" if id3.get("TSRC") is None else str(id3["TSRC"])
    id3value.encodedby = "" if id3.get("TENC") is None else str(id3["TENC"])
    id3value.size = int(str(id3["TSIZ"]))
    id3value.header_size = id3.size
    id3value.encoding_technology = str(id3["TSSE"])
    return id3value

def get_printable_count(x:bytes):
    i = 0
    for i,c in enumerate(x):
        # all pritable
        if c < 0x20 or c > 0x7e:
            return i
    return i

def get_printable_bytes(x:bytes):
    return x[:get_printable_count(x)]

def xm_decrypt(raw_data):
    # decode id3
    xm_info = get_xm_info(raw_data)
    print("id3 header size: ",hex(xm_info.header_size))
    encrypted_data = raw_data[xm_info.header_size:xm_info.header_size+xm_info.size:]

    # Stage 1 aes-256-cbc
    xm_key = b"ximalayaximalayaximalayaximalaya"
    print(f"decrypt stage 1 (aes-256-cbc):\n"
          f"    data length = {len(encrypted_data)},\n"
          f"    key = {xm_key},\n"
          f"    iv = {xm_info.iv().hex()}")
    cipher = AES.new(xm_key, AES.MODE_CBC, xm_info.iv())
    de_data = cipher.decrypt(pad(encrypted_data, 16))
    # Stage 2 xmDecrypt = (base64 decode => aes-192-cbc => base64 encode)
    print(f"decrypt stage 2 (xmDecrypt):\n"
          f"    data length = {len(de_data)},\n"
          f"    key = {str(xm_info.tracknumber)}")
    stage_2_data = base64.b64decode(get_printable_bytes(de_data))
    assert len(stage_2_data) % 16 == 0
    key = str(xm_info.tracknumber).encode()
    key = (b'12345678'*3)[:0x18-len(key)] + key
    cipher = AES.new(key, AES.MODE_CBC, key[:16])
    stage_2_data = unpad(cipher.decrypt(stage_2_data),16).decode() # idk but workround
    # Stage 3 combine
    print(f"Stage 3 (base64 combination):\n"
          f"    technology = {xm_info.encoding_technology}")
    decrypted_data = base64.b64decode(xm_info.encoding_technology+stage_2_data)
    final_data = decrypted_data + raw_data[xm_info.header_size+xm_info.size::]
    return xm_info,final_data

def xm_decrypt_v12():
    pass

def find_ext(data):
    exts = ["m4a","mp3","flac","wav"]
    value = magic.from_buffer(data).lower()
    for ext in exts:
        if ext in value:
            return ext
    raise Exception(f"unexpected format {value}")

def decrypt_xm_file(from_file,output=''):
    print(f"decrypting {from_file}")
    data = read_file(from_file)
    info, audio_data = xm_decrypt(data)
    if output == "":
        output = re.sub(r'[^\w\-_\. ]', '_', info.title)+"."+find_ext(audio_data[:0xff])
    buffer = io.BytesIO(audio_data)
    tags = mutagen.File(buffer,easy=True)
    tags["title"] = info.title
    tags["album"] = info.album
    tags["artist"] = info.artist
    print(tags.pprint())
    tags.save(buffer)
    with open(output,"wb") as f:
        buffer.seek(0)
        f.write(buffer.read())
    print(f"decrypt succeed, file write to {output}")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python decrypt_xm.py [<filename> ...]")
    for filename in sys.argv[1::]:
        decrypt_xm_file(filename)

def iv(self):
    if self.ISRC != "":
        return bytes.fromhex(self.ISRC)
    return bytes.fromhex(self.encodedby)

# Stage 1 aes-256-cbc
xm_key = b"ximalayaximalayaximalayaximalaya"
# print(f"decrypt stage 1 (aes-256-cbc):\n"
#       f"    data length = {len(encrypted_data)},\n"
#       f"    key = {xm_key},\n"
#       f"    iv = {xm_info.iv().hex()}")
cipher = AES.new(xm_key, AES.MODE_CBC, xm_info.iv())
de_data = cipher.decrypt(pad(encrypted_data, 16))
# print("success")
# Stage 2 xmDecrypt
de_data = get_printable_bytes(de_data)
track_id = str(xm_info.tracknumber).encode()
stack_pointer = xm_encryptor.exports.a(-16)
assert isinstance(stack_pointer, int)
de_data_offset = xm_encryptor.exports.c(len(de_data))
assert isinstance(de_data_offset, int)
track_id_offset = xm_encryptor.exports.c(len(track_id))
assert isinstance(track_id_offset, int)
memory_i = xm_encryptor.exports.i
memview_unit8: Uint8Array = memory_i.uint8_view(offset=de_data_offset)
for i, b in enumerate(de_data):
    memview_unit8[i] = b
memview_unit8: Uint8Array = memory_i.uint8_view(offset=track_id_offset)
for i, b in enumerate(track_id):
    memview_unit8[i] = b
# print(bytearray(memory_i.buffer)[track_id_offset:track_id_offset + len(track_id)].decode())
# print(f"decrypt stage 2 (xmDecrypt):\n"
#       f"    stack_pointer = {stack_pointer},\n"
#       f"    data_pointer = {de_data_offset}, data_length = {len(de_data)},\n"
#       f"    track_id_pointer = {track_id_offset}, track_id_length = {len(track_id)}")
# print("success")
xm_encryptor.exports.g(stack_pointer, de_data_offset, len(de_data), track_id_offset, len(track_id))
memview_int32: Int32Array = memory_i.int32_view(offset=stack_pointer // 4)
result_pointer = memview_int32[0]
result_length = memview_int32[1]
assert memview_int32[2] == 0, memview_int32[3] == 0
result_data = bytearray(memory_i.buffer)[result_pointer:result_pointer + result_length].decode()
# Stage 3 combine
# print(f"Stage 3 (base64)")
decrypted_data = base64.b64decode(xm_info.encoding_technology + result_data)
final_data = decrypted_data + raw_data[xm_info.header_size + xm_info.size::]
# print("success")
return xm_info, final_data

album_path = f"{output_path}/{replace_invalid_chars(info.album)}"
temp_output = f"{album_path}/{replace_invalid_chars(info.title)}.aac"

if not os.path.exists(album_path):
    os.makedirs(album_path)

buffer = io.BytesIO(audio_data)
tags = mutagen.File(buffer, easy=True)
tags["title"] = info.title
tags["album"] = info.album
tags["artist"] = info.artist
tags.save(buffer)

with open(temp_output, "wb") as f:
    buffer.seek(0)
    f.write(buffer.read())

return temp_output

for xm_file in xm_files:
    print(f"正在解密文件: {xm_file}")
    aac_file = decrypt_xm_file_to_aac(xm_file, output_path)
    aac_files.append(aac_file)

with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
    futures = []
    for aac_file in aac_files:
        mp3_file = aac_file.replace(".aac", ".mp3")
        futures.append(executor.submit(convert_to_mp3_in_thread, aac_file, mp3_file, '24k'))

    for future in concurrent.futures.as_completed(futures):
        print("文件转换完成:", future.result())

import base64
import concurrent.futures
import glob
import io
import os
import pathlib
import subprocess
import sys
import mutagen
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from mutagen.easyid3 import ID3
from wasmer import Store, Module, Instance, Uint8Array, Int32Array, engine
from wasmer_compiler_cranelift import Compiler

class XMInfo:
    def __init__(self):
        self.title = ""
        self.artist = ""
        self.album = ""
        self.tracknumber = 0
        self.size = 0
        self.header_size = 0
        self.ISRC = ""
        self.encodedby = ""
        self.encoding_technology = ""

    def iv(self):
        if self.ISRC != "":
            return bytes.fromhex(self.ISRC)
        return bytes.fromhex(self.encodedby)


def get_str(x):
    if x is None:
        return ""
    return x


def read_file(x):
    with open(x, "rb") as f:
        return f.read()


# return number of id3 bytes
def get_xm_info(data: bytes):
    # print(EasyID3(io.BytesIO(data)))
    id3 = ID3(io.BytesIO(data), v2_version=3)
    id3value = XMInfo()
    id3value.title = str(id3["TIT2"])
    id3value.album = str(id3["TALB"])
    id3value.artist = str(id3["TPE1"])
    id3value.tracknumber = int(str(id3["TRCK"]))
    id3value.ISRC = "" if id3.get("TSRC") is None else str(id3["TSRC"])
    id3value.encodedby = "" if id3.get("TENC") is None else str(id3["TENC"])
    id3value.size = int(str(id3["TSIZ"]))
    id3value.header_size = id3.size
    id3value.encoding_technology = str(id3["TSSE"])
    return id3value


def get_printable_count(x: bytes):
    i = 0
    for i, c in enumerate(x):
        # all pritable
        if c < 0x20 or c > 0x7e:
            return i
    return i


def get_printable_bytes(x: bytes):
    return x[:get_printable_count(x)]

def xm_decrypt(raw_data):
    xm_encryptor = Instance(Module(
        Store(engine.Universal(Compiler)),
        pathlib.Path("./xm_encryptor.wasm").read_bytes()
    ))
    xm_info = get_xm_info(raw_data)
    encrypted_data = raw_data[xm_info.header_size:xm_info.header_size + xm_info.size:]
    xm_key = b"ximalayaximalayaximalayaximalaya"
    cipher = AES.new(xm_key, AES.MODE_CBC, xm_info.iv())
    de_data = cipher.decrypt(pad(encrypted_data, 16))
    de_data = get_printable_bytes(de_data)
    track_id = str(xm_info.tracknumber).encode()
    stack_pointer = xm_encryptor.exports.a(-16)
    assert isinstance(stack_pointer, int)
    de_data_offset = xm_encryptor.exports.c(len(de_data))
    assert isinstance(de_data_offset, int)
    track_id_offset = xm_encryptor.exports.c(len(track_id))
    assert isinstance(track_id_offset, int)
    memory_i = xm_encryptor.exports.i
    memview_unit8: Uint8Array = memory_i.uint8_view(offset=de_data_offset)
    for i, b in enumerate(de_data):
        memview_unit8[i] = b
    memview_unit8: Uint8Array = memory_i.uint8_view(offset=track_id_offset)
    for i, b in enumerate(track_id):
        memview_unit8[i] = b
    xm_encryptor.exports.g(stack_pointer, de_data_offset, len(de_data), track_id_offset, len(track_id))
    memview_int32: Int32Array = memory_i.int32_view(offset=stack_pointer // 4)
    result_pointer = memview_int32[0]
    result_length = memview_int32[1]
    assert memview_int32[2] == 0, memview_int32[3] == 0
    result_data = bytearray(memory_i.buffer)[result_pointer:result_pointer + result_length].decode()
    decrypted_data = base64.b64decode(xm_info.encoding_technology + result_data)
    final_data = decrypted_data + raw_data[xm_info.header_size + xm_info.size::]
    return xm_info, final_data

def replace_invalid_chars(name):
    invalid_chars = ['/', '\\', ':', '*', '?', '"', '<', '>', '|']
    for char in invalid_chars:
        if char in name:
            name = name.replace(char, " ")
    return name

def convert_to_mp3(input_file, output_file, bitrate='24k'):
    command = ['ffmpeg', '-i', input_file, '-acodec', 'libmp3lame', '-b:a', bitrate, output_file]
    subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

def decrypt_xm_file_to_aac(from_file, output_path='./output'):
    data = read_file(from_file)
    info, audio_data = xm_decrypt(data)

    album_path = f"{output_path}/{replace_invalid_chars(info.album)}"
    temp_output = f"{album_path}/{replace_invalid_chars(info.title)}.aac"

    if not os.path.exists(album_path):
        os.makedirs(album_path)

    buffer = io.BytesIO(audio_data)
    tags = mutagen.File(buffer, easy=True)
    tags["title"] = info.title
    tags["album"] = info.album
    tags["artist"] = info.artist
    tags.save(buffer)

    with open(temp_output, "wb") as f:
        buffer.seek(0)
        f.write(buffer.read())

    return temp_output

def convert_to_mp3_in_thread(aac_file, mp3_file, bitrate='24k'):
    convert_to_mp3(aac_file, mp3_file, bitrate)
    os.remove(aac_file)  # 删除临时AAC文件

def batch_decrypt_and_convert(directory, output_path='./output'):
    xm_files = glob.glob(os.path.join(directory, "*.xm"))
    aac_files = []

    for xm_file in xm_files:
        print(f"正在解密文件: {xm_file}")
        aac_file = decrypt_xm_file_to_aac(xm_file, output_path)
        aac_files.append(aac_file)

    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
        futures = []
        for aac_file in aac_files:
            mp3_file = aac_file.replace(".aac", ".mp3")
            futures.append(executor.submit(convert_to_mp3_in_thread, aac_file, mp3_file, '24k'))

        for future in concurrent.futures.as_completed(futures):
            print("文件转换完成:", future.result())


def get_output_path():
    print("请选择是否需要设置输出路径：（不设置默认为本程序目录下的output文件夹）")
    print("1. 设置输出路径")
    print("2. 不设置输出路径")
    choice = input()
    if choice == "1":
        print("请输入输出路径：")
        output_path = input()
        if not os.path.exists(output_path):
            os.makedirs(output_path)
        return output_path
    else:
        return "./output"

if __name__ == "__main__":
    while True:
        print("欢迎使用喜马拉雅音频解密工具")
        print("本工具仅供学习交流使用，严禁用于商业用途")
        print("请选择您想要使用的功能：")
        print("1. 解密单个文件")
        print("2. 批量解密文件")
        print("3. 退出")
        choice = input()
        if choice == "1":
            print("请输入需要解密的文件路径：")
            file_to_decrypt = input()
            if os.path.exists(file_to_decrypt) and os.path.isfile(file_to_decrypt):
                output_path = get_output_path()
                decrypt_xm_file_to_aac(file_to_decrypt, output_path)
            else:
                print("您输入的不是一个有效的文件路径，请重新输入！")
        elif choice == "2":
            print("请输入包含需要解密的文件的文件夹路径：")
            dir_to_decrypt = input()
            if os.path.exists(dir_to_decrypt) and os.path.isdir(dir_to_decrypt):
                output_path = get_output_path()
                batch_decrypt_and_convert(dir_to_decrypt, output_path)
            else:
                print("您输入的不是一个有效的文件夹路径，请重新输入！")
        elif choice == "3":
            print("退出程序。")
            sys.exit()
        else:
            print("输入错误，请重新输入！")