Add WavPack MD5 metadata support

mutagen/wavpack.py

@@ -12,12 +12,15 @@
 
 * http://www.wavpack.com/
 * http://www.wavpack.com/file_format.txt
+* http://www.wavpack.com/WavPack5FileFormat.pdf
 
 for more information.
 """
 
 __all__ = ["WavPack", "Open", "delete"]
 
+from functools import reduce
+
 from mutagen import StreamInfo
 from mutagen.apev2 import APEv2File, error, delete
 from mutagen._util import cdata, convert_error
@@ -29,6 +32,20 @@ class WavPackHeaderError(error):
 RATES = [6000, 8000, 9600, 11025, 12000, 16000, 22050, 24000, 32000, 44100,
          48000, 64000, 88200, 96000, 192000]
 
+# Metadata sub-block ID flags (from wavpack.h)
+_ID_LARGE = 0x80         # size field is 3 bytes instead of 1
+_ID_ODD_SIZE = 0x40      # actual data is 1 byte less than size_words*2
+_ID_UNIQUE = 0x3f        # mask for the unique sub-block identifier
+
+# MD5 sub-block identifiers (ID_OPTIONAL_DATA | n, see wavpack.h)
+_ID_MD5_CHECKSUM = 0x26      # standard PCM files
+_ID_ALT_MD5_CHECKSUM = 0x29  # DSD / alternate-format files (WavPack 5)
+
+
+def _to_int_be(data):
+    """Convert a byte string to a long using big-endian byte order."""
+    return reduce(lambda a, b: (a << 8) + b, bytearray(data), 0)
+
 
 class _WavPackHeader(object):
 
@@ -70,6 +87,93 @@ def from_fileobj(cls, fileobj):
                               samples, block_index, block_samples, flags, crc)
 
 
+def _extract_md5_from_metadata(fileobj, block_size):
+    """Extract MD5 signature from WavPack metadata sub-blocks.
+
+    Parses the metadata sub-blocks that follow the 32-byte frame header in
+    the initial audio block, looking for the MD5 checksum sub-block
+    (ID 0x26 for standard PCM files, ID 0x29 for DSD / alternate-format
+    files). Returns None if no MD5 sub-block is present or if parsing fails.
+
+    Each sub-block is laid out as follows (from wavpack.h):
+
+      - 1 byte : ID byte
+          bit 7 (_ID_LARGE)    - if set, size field is 3 bytes; else 1 byte
+          bit 6 (_ID_ODD_SIZE) - if set, actual data is size_words*2 - 1 bytes
+          bits 5-0             - unique sub-block identifier (_ID_UNIQUE mask)
+      - 1 or 3 bytes : size in 16-bit words (little-endian for 3-byte form)
+      - size_words*2 bytes : data payload (last byte is padding if _ID_ODD_SIZE)
+
+    Args:
+        fileobj: file-like object positioned immediately after the 32-byte
+            frame header
+        block_size (int): ckSize field from the frame header, which counts
+            bytes from offset 8 to the end of the block; the metadata
+            payload remaining after the 32-byte header is block_size - 24
+
+    Returns:
+        int or None: MD5 signature as a big-endian integer, or None if the
+            sub-block is absent or the data is malformed
+    """
+    # The ckSize field covers everything after the first 8 bytes of the block.
+    # We have already consumed the full 32-byte header, so the remaining
+    # metadata payload is:  ckSize - (32 - 8) = block_size - 24 bytes.
+    # (This matches the seek expression used elsewhere: block_size - 32 + 8.)
+    metadata_size = block_size - 24
+    if metadata_size <= 0:
+        return None
+
+    try:
+        metadata_data = fileobj.read(metadata_size)
+        if len(metadata_data) != metadata_size:
+            return None
+    except IOError:
+        return None
+
+    offset = 0
+    while offset < len(metadata_data):
+        if offset >= len(metadata_data):
+            break
+
+        id_byte = metadata_data[offset]
+        offset += 1
+
+        is_large = bool(id_byte & _ID_LARGE)
+        is_odd = bool(id_byte & _ID_ODD_SIZE)
+        unique_id = id_byte & _ID_UNIQUE
+
+        # Size is stored in 16-bit words (not bytes).
+        # _ID_LARGE in the ID byte selects 3-byte vs 1-byte size encoding.
+        if is_large:
+            if offset + 3 > len(metadata_data):
+                break
+            word_count = (metadata_data[offset] |
+                          (metadata_data[offset + 1] << 8) |
+                          (metadata_data[offset + 2] << 16))
+            offset += 3
+        else:
+            if offset >= len(metadata_data):
+                break
+            word_count = metadata_data[offset]
+            offset += 1
+
+        # The block always occupies word_count*2 bytes on disk; _ID_ODD_SIZE
+        # means the last byte of that allocation is padding, not payload.
+        data_size = word_count * 2 - (1 if is_odd else 0)
+
+        if unique_id in (_ID_MD5_CHECKSUM, _ID_ALT_MD5_CHECKSUM):
+            if data_size == 16:
+                md5_data = metadata_data[offset:offset + 16]
+                if len(md5_data) == 16:
+                    return _to_int_be(md5_data)
+            return None
+
+        # Advance past the full word-aligned block payload
+        offset += word_count * 2
+
+    return None
+
+
 class WavPackInfo(StreamInfo):
     """WavPack stream information.
 
@@ -79,8 +183,12 @@ class WavPackInfo(StreamInfo):
         sample_rate (int): audio sampling rate in Hz
         bits_per_sample (int): audio sample size
         version (int): WavPack stream version
+        md5_signature (int or None): MD5 checksum of the original
+            uncompressed audio as an integer, or None if not present.
     """
 
+    md5_signature = None
+
     def __init__(self, fileobj):
         try:
             header = _WavPackHeader.from_fileobj(fileobj)
@@ -97,21 +205,33 @@ def __init__(self, fileobj):
             self.sample_rate *= 4
             self.bits_per_sample = 1
 
-        if header.total_samples == -1 or header.block_index != 0:
-            # TODO: we could make this faster by using the tag size
-            # and search backwards for the last block, then do
-            # last.block_index + last.block_samples - initial.block_index
+        # Parse metadata from the first block.
+        self.md5_signature = _extract_md5_from_metadata(fileobj,
+                                                        header.block_size)
+
+        need_samples = header.total_samples == -1 or header.block_index != 0
+        if need_samples:
             samples = header.block_samples
-            while 1:
-                fileobj.seek(header.block_size - 32 + 8, 1)
-                try:
-                    header = _WavPackHeader.from_fileobj(fileobj)
-                except WavPackHeaderError:
-                    break
-                samples += header.block_samples
         else:
             samples = header.total_samples
 
+        # Continue scanning blocks if we still need stream length or if MD5
+        # wasn't present in the first block.
+        while need_samples or self.md5_signature is None:
+            try:
+                header = _WavPackHeader.from_fileobj(fileobj)
+            except WavPackHeaderError:
+                break
+
+            if self.md5_signature is None:
+                self.md5_signature = _extract_md5_from_metadata(
+                    fileobj, header.block_size)
+            else:
+                fileobj.seek(header.block_size - 32 + 8, 1)
+
+            if need_samples:
+                samples += header.block_samples
+
         self.length = float(samples) / self.sample_rate
 
     def pprint(self):

tests/test_wavpack.py

@@ -1,7 +1,9 @@
-
 import os
+from io import BytesIO
 
-from mutagen.wavpack import WavPack, WavPackHeaderError
+from mutagen.wavpack import (WavPack, WavPackHeaderError,
+                             _to_int_be, _extract_md5_from_metadata,
+                             _ID_MD5_CHECKSUM, _ID_ALT_MD5_CHECKSUM)
 from tests import TestCase, DATA_DIR
 
 
@@ -25,6 +27,14 @@ def test_bits_per_sample(self):
     def test_length(self):
         self.failUnlessAlmostEqual(self.audio.info.length, 3.68, 2)
 
+    def test_md5_signature(self):
+        # The test file may or may not carry an MD5 (depends on how it was
+        # encoded).  Just verify the attribute has the correct type.
+        self.assertTrue(
+            self.audio.info.md5_signature is None or
+            isinstance(self.audio.info.md5_signature, int)
+        )
+
     def test_not_my_file(self):
         self.failUnlessRaises(
             WavPackHeaderError, WavPack, os.path.join(DATA_DIR, "empty.ogg"))
@@ -88,3 +98,121 @@ def test_pprint(self):
 
     def test_mime(self):
         self.failUnless("audio/x-wavpack" in self.audio.mime)
+
+
+class TMD5Extraction(TestCase):
+    """Unit tests for WavPack metadata sub-block MD5 parsing.
+
+    Sub-block wire format (from wavpack.h):
+      - 1 byte ID: bit7=LARGE, bit6=ODD_SIZE, bits5-0=unique id
+      - 1 byte (or 3 bytes if LARGE) size in 16-bit words
+      - size_words*2 bytes of payload data
+    """
+
+    def _make_fileobj(self, data):
+        return BytesIO(data)
+
+    # ------------------------------------------------------------------
+    # _to_int_be
+    # ------------------------------------------------------------------
+
+    def test_to_int_be_zero(self):
+        self.failUnlessEqual(_to_int_be(b'\x00' * 16), 0)
+
+    def test_to_int_be_one(self):
+        self.failUnlessEqual(_to_int_be(b'\x00' * 15 + b'\x01'), 1)
+
+    def test_to_int_be_msb(self):
+        self.failUnlessEqual(_to_int_be(b'\x01' + b'\x00' * 15),
+                             1 << 120)
+
+    def test_to_int_be_all_ff(self):
+        self.failUnlessEqual(_to_int_be(b'\xff' * 4), 0xffffffff)
+
+    # ------------------------------------------------------------------
+    # _extract_md5_from_metadata: boundary / empty cases
+    # ------------------------------------------------------------------
+
+    def test_no_metadata_when_block_size_too_small(self):
+        # block_size - 24 <= 0  →  no metadata possible
+        result = _extract_md5_from_metadata(BytesIO(b''), 24)
+        self.failUnlessEqual(result, None)
+
+    def test_empty_payload(self):
+        result = _extract_md5_from_metadata(BytesIO(b''), 25)
+        self.failUnlessEqual(result, None)
+
+    # ------------------------------------------------------------------
+    # Correct ID / correct size → should find MD5
+    # ------------------------------------------------------------------
+
+    def test_standard_md5_single_byte_size(self):
+        # ID_MD5_CHECKSUM = 0x26; 16 bytes = 8 words → size byte = 0x08
+        md5_bytes = bytes(range(16))
+        payload = bytes([_ID_MD5_CHECKSUM, 0x08]) + md5_bytes
+        result = _extract_md5_from_metadata(BytesIO(payload),
+                                            24 + len(payload))
+        self.failUnlessEqual(result, _to_int_be(md5_bytes))
+
+    def test_alt_md5_single_byte_size(self):
+        # ID_ALT_MD5_CHECKSUM = 0x29 (used for DSD / alt-format files)
+        md5_bytes = bytes(range(16, 32))
+        payload = bytes([_ID_ALT_MD5_CHECKSUM, 0x08]) + md5_bytes
+        result = _extract_md5_from_metadata(BytesIO(payload),
+                                            24 + len(payload))
+        self.failUnlessEqual(result, _to_int_be(md5_bytes))
+
+    def test_md5_found_after_other_sub_blocks(self):
+        # A realistic payload: one non-MD5 sub-block followed by the MD5.
+        # Sub-block 1: ID=0x01 (ID_ENCODER_INFO), 2 words (4 bytes) of data
+        other = bytes([0x01, 0x02]) + b'\x00' * 4
+        md5_bytes = b'\xde\xad\xbe\xef' * 4
+        md5_block = bytes([_ID_MD5_CHECKSUM, 0x08]) + md5_bytes
+        payload = other + md5_block
+        result = _extract_md5_from_metadata(BytesIO(payload),
+                                            24 + len(payload))
+        self.failUnlessEqual(result, _to_int_be(md5_bytes))
+
+    # ------------------------------------------------------------------
+    # Wrong ID or wrong size → should return None
+    # ------------------------------------------------------------------
+
+    def test_no_md5_block_present(self):
+        # Only a non-MD5 sub-block; result must be None
+        payload = bytes([0x01, 0x02]) + b'\x00' * 4
+        result = _extract_md5_from_metadata(BytesIO(payload),
+                                            24 + len(payload))
+        self.failUnlessEqual(result, None)
+
+    def test_wrong_size_returns_none(self):
+        # MD5 block claims 15 bytes (7 words + ODD_SIZE), not 16 → rejected
+        payload = bytes([_ID_MD5_CHECKSUM | 0x40, 0x08]) + b'\x00' * 16
+        result = _extract_md5_from_metadata(BytesIO(payload),
+                                            24 + len(payload))
+        self.failUnlessEqual(result, None)
+
+    def test_truncated_md5_data_returns_none(self):
+        # Header says 8 words but the fileobj only has 15 bytes of payload
+        payload = bytes([_ID_MD5_CHECKSUM, 0x08]) + b'\x00' * 15
+        result = _extract_md5_from_metadata(BytesIO(payload),
+                                            24 + len(payload))
+        self.failUnlessEqual(result, None)
+
+    # ------------------------------------------------------------------
+    # Large (3-byte) size encoding
+    # ------------------------------------------------------------------
+
+    def test_large_size_non_md5_block_skipped(self):
+        # Build a large sub-block (ID_LARGE set) with 200 bytes of data
+        # (100 words), followed by the MD5 block.
+        large_data = b'\xab' * 200
+        # ID byte with _ID_LARGE set, arbitrary unique id 0x05
+        large_id = 0x80 | 0x05
+        # 3-byte little-endian word count: 100 = 0x64
+        large_block = bytes([large_id, 100, 0, 0]) + large_data
+        md5_bytes = b'\x11\x22\x33\x44' * 4
+        md5_block = bytes([_ID_MD5_CHECKSUM, 0x08]) + md5_bytes
+        payload = large_block + md5_block
+        result = _extract_md5_from_metadata(BytesIO(payload),
+                                            24 + len(payload))
+        self.failUnlessEqual(result, _to_int_be(md5_bytes))