# -*- coding: utf-8 -*-
"""网易云音乐 .ncm 解密模块 v1

纯标准库实现(零依赖):AES-128-ECB 仅用于解密几百字节的密钥块和元数据;
音频为 256 字节 RC4 式密钥盒循环 XOR,已装 numpy 时自动向量化加速(98MB 约 0.1 秒),
未装则用预计算转换表回退(纯 Python 约 10~30 MB/s)。

算法已与开源实现(ncmdump-py / UnlockMusic)端到端验证,4 个测试文件字节级一致。

用法:
    import ncm_decrypt
    ncm_decrypt.decrypt_file('xxx.ncm', 'xxx.flac')          # 解密音频
    key, meta, cover, audio_off = ncm_decrypt.parse_ncm(...) # 只解析头/元数据
"""

import base64
import json
import os
import struct

__all__ = ['parse_ncm', 'decrypt_audio', 'decrypt_file']

CORE_KEY = bytes.fromhex('687a4852416d736f356b496e62617857')   # 'hzHRAmso5kInbaxW'
META_KEY = bytes.fromhex('2331346C6A6B5F215C5D2630553C2728')   # '#14ljk_!\]&0U<\'('
MAGIC = b'CTENFDAM'

# ---------- 微型 AES-128(仅解密方向,ECB) ----------
# 逆 S 盒由正 S 盒程序化推导,杜绝手打笔误
_SBOX = (
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16,
)
_INV_SBOX = [0] * 256
for _i, _v in enumerate(_SBOX):
    _INV_SBOX[_v] = _i
_RCON = (0x00, 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36)
_ROUNDS = 10


def _bytes2matrix(text):
    return [list(text[i:i + 4]) for i in range(0, 16, 4)]


def _matrix2bytes(matrix):
    return bytes(sum(matrix, []))


def _xtime(a):
    return (((a << 1) ^ 0x1B) & 0xFF) if (a & 0x80) else (a << 1)


def _add_round_key(s, k):
    for i in range(4):
        for j in range(4):
            s[i][j] ^= k[i][j]


def _expand_key(master_key):
    key_columns = _bytes2matrix(master_key)
    i = 1
    while len(key_columns) < (_ROUNDS + 1) * 4:
        word = list(key_columns[-1])
        if len(key_columns) % 4 == 0:
            word.append(word.pop(0))
            word = [_SBOX[b] for b in word]
            word[0] ^= _RCON[i]
            i += 1
        word = bytes(p ^ q for p, q in zip(word, key_columns[-4]))
        key_columns.append(word)
    return [key_columns[4 * i:4 * (i + 1)] for i in range(len(key_columns) // 4)]


def _inv_shift_rows(s):
    s[0][1], s[1][1], s[2][1], s[3][1] = s[3][1], s[0][1], s[1][1], s[2][1]
    s[0][2], s[1][2], s[2][2], s[3][2] = s[2][2], s[3][2], s[0][2], s[1][2]
    s[0][3], s[1][3], s[2][3], s[3][3] = s[1][3], s[2][3], s[3][3], s[0][3]


def _inv_sub_bytes(s):
    for i in range(4):
        for j in range(4):
            s[i][j] = _INV_SBOX[s[i][j]]


def _mix_columns(s):
    for i in range(4):
        t = s[i][0] ^ s[i][1] ^ s[i][2] ^ s[i][3]
        u = s[i][0]
        s[i][0] ^= t ^ _xtime(s[i][0] ^ s[i][1])
        s[i][1] ^= t ^ _xtime(s[i][1] ^ s[i][2])
        s[i][2] ^= t ^ _xtime(s[i][2] ^ s[i][3])
        s[i][3] ^= t ^ _xtime(s[i][3] ^ u)


def _inv_mix_columns(s):
    for i in range(4):
        u = _xtime(_xtime(s[i][0] ^ s[i][2]))
        v = _xtime(_xtime(s[i][1] ^ s[i][3]))
        s[i][0] ^= u
        s[i][1] ^= v
        s[i][2] ^= u
        s[i][3] ^= v
    _mix_columns(s)


def _aes_ecb_decrypt(data, key16):
    rk = _expand_key(key16)
    out = bytearray()
    for i in range(0, len(data), 16):
        s = _bytes2matrix(data[i:i + 16])
        _add_round_key(s, rk[-1])
        for rnd in range(_ROUNDS - 1, 0, -1):
            _inv_shift_rows(s)
            _inv_sub_bytes(s)
            _add_round_key(s, rk[rnd])
            _inv_mix_columns(s)
        _inv_shift_rows(s)
        _inv_sub_bytes(s)
        _add_round_key(s, rk[0])
        out += _matrix2bytes(s)
    return bytes(out)


def _pkcs7_unpad(data):
    """去掉 PKCS#7 填充(填充长度 = 末尾字节值)。"""
    n = data[-1]
    if not (1 <= n <= 16 and data[-n:] == bytes([n]) * n):
        raise ValueError('PKCS#7 填充无效')
    return data[:-n]


# ---------- ncm 解析与解密 ----------

def _keybox(key):
    """RC4 式密钥盒: 标准 KSA + 一轮 PRGA 生成 256 字节循环密钥流。"""
    S = list(range(256))
    j = 0
    for a in range(256):
        j = (S[a] + j + key[a % len(key)]) & 255
        S[a], S[j] = S[j], S[a]
    ks = [0] * 256
    for k in range(256):
        a = S[(k + 1) & 255]
        b = S[((k + 1) + a) & 255]
        ks[k] = S[(a + b) & 255]
    return ks


def parse_ncm(path):
    """解析 .ncm 文件头。

    返回 (rc4_key, meta_dict, cover_bytes, audio_offset)。
    meta_dict 含 musicName/artist/album/format 等字段(元数据为空时为 {})。
    """
    with open(path, 'rb') as f:
        b = f.read()
    if b[:8] != MAGIC:
        raise ValueError('不是有效的 ncm 文件(魔数不匹配)')

    # 密钥块: XOR 0x64 -> AES-128-ECB(Core Key) -> 去填充 -> 去 17 字节前缀
    key_len = struct.unpack('<I', b[10:14])[0]
    key_block = bytes(c ^ 0x64 for c in b[14:14 + key_len])
    dec = _pkcs7_unpad(_aes_ecb_decrypt(key_block, CORE_KEY))
    if not dec.startswith(b'neteasecloudmusic'):
        raise ValueError('密钥块解密失败(前缀不匹配)')
    rc4_key = dec[len(b'neteasecloudmusic'):]  # 17 字节

    # 元数据: XOR 0x63 -> 去 22 字节前缀 -> base64 -> AES-128-ECB(Meta Key)
    meta_off = 14 + key_len
    meta_len = struct.unpack('<I', b[meta_off:meta_off + 4])[0]
    meta = {}
    if meta_len:
        m_block = bytes(c ^ 0x63 for c in b[meta_off + 4:meta_off + 4 + meta_len])
        m_dec = _pkcs7_unpad(_aes_ecb_decrypt(base64.b64decode(m_block[22:]), META_KEY))
        i = m_dec.index(b':')
        data = json.loads(m_dec[i + 1:].decode('utf-8', 'replace'))
        meta = data['mainMusic'] if m_dec[:i] == b'dj' else data  # dj 电台变体取 mainMusic

    # 音频起点: CRC(4) + 未知(5) + 封面长度(4) + 封面(未加密)
    meta_end = meta_off + 4 + meta_len
    cover_len = struct.unpack('<I', b[meta_end + 9:meta_end + 13])[0]
    cover = b[meta_end + 13:meta_end + 13 + cover_len]
    audio_off = meta_end + 13 + cover_len
    return rc4_key, meta, cover, audio_off


def decrypt_audio(path, out_path):
    """解密 .ncm 的音频部分写入 out_path,返回 meta_dict。

    numpy 已装时向量化解密(百 MB 秒级),否则用预计算转换表回退。
    """
    rc4_key, meta, _, audio_off = parse_ncm(path)
    ks = _keybox(rc4_key)
    with open(path, 'rb') as f:
        f.seek(audio_off)
        enc = f.read()
    n = len(enc)
    try:  # numpy 加速路径
        import numpy as np
        kst = np.tile(np.frombuffer(bytes(ks), dtype=np.uint8), (n + 255) // 256)[:n]
        out = (np.frombuffer(enc, dtype=np.uint8) ^ kst).tobytes()
    except ImportError:  # 纯标准库回退: 预计算 256 张 XOR 转换表
        tables = [bytes(x ^ k for x in range(256)) for k in ks]
        out = bytearray()
        for j in range(0, n, 1 << 20):  # 按 1MB 分块,块内每 256 字节轮换表
            chunk = enc[j:j + (1 << 20)]
            out += b''.join(chunk[t:t + 256].translate(tables[(j // 256 + t // 256) & 255])
                            for t in range(0, len(chunk), 256))
        out = bytes(out)
    with open(out_path, 'wb') as f:
        f.write(out)
    return meta


def decrypt_file(src_path, out_path=None):
    """解密整个 .ncm 文件。out_path 缺省时输出到同目录同名 .flac/.mp3(按元数据格式)。"""
    meta = parse_ncm(src_path)[1]
    if out_path is None:
        ext = '.' + (meta.get('format') or 'mp3')
        out_path = os.path.splitext(src_path)[0] + ext
    decrypt_audio(src_path, out_path)
    return out_path


if __name__ == '__main__':
    # 自检: NIST SP 800-38A F.1.1 标准向量(AES-128-ECB 解密)
    k = bytes(range(16))
    ct = bytes.fromhex('69c4e0d86a7b0430d8cdb78070b4c55a')
    ok = _aes_ecb_decrypt(ct, k).hex() == '00112233445566778899aabbccddeeff'
    print('AES 自检:', '通过' if ok else '失败!')
    sys_exit = 0 if ok else 1
    import sys
    sys.exit(sys_exit)