diora-web/podcasts/id3.py

"""Chapters embedded in the ID3v2 tag at the start of an MP3 file.

Written from the ID3v2.3.0 and ID3v2.4.0 specifications and the ID3v2
Chapter Frame Addendum (CHAP and CTOC frames). Only what chapters need is
read: the tag header, an optional extended header (skipped), and the frames.
Of those, CHAP frames give the chapters, with the title from an embedded TIT2
frame and a link from WXXX or another W*** frame; CTOC frames, when there
are any, give the order (which matters only for chapters starting at the
same time, as callers sort by start). Pictures (APIC) and everything else are
skipped.

Pure parsing, no I/O: the caller hands over the first bytes of the file
(`tag_size` says how many are needed).
"""

import re
import struct
import zlib

HEADER_SIZE = 10
MAX_FRAME_BYTES = 4 * 1024 * 1024   # a decompressed frame can't grow past this

_FRAME_ID = re.compile(rb'[A-Z0-9]{4}\Z')

# Header flags
_UNSYNC = 0x80
_EXTENDED = 0x40
_FOOTER = 0x10


def _synchsafe(raw):
    """Integer from bytes that carry 7 bits each (the high bit is always 0)."""
    value = 0
    for b in raw:
        value = (value << 7) | (b & 0x7F)
    return value


def tag_size(header):
    """Bytes the ID3v2 tag at the start of a file takes up, header and footer
    included, from the file's first 10 bytes; None if there's no such tag."""
    if len(header) < HEADER_SIZE or header[:3] != b'ID3':
        return None
    major, revision, flags = header[3], header[4], header[5]
    if major == 0xFF or revision == 0xFF or any(b & 0x80 for b in header[6:10]):
        return None
    size = HEADER_SIZE + _synchsafe(header[6:10])
    if major >= 4 and flags & _FOOTER:
        size += HEADER_SIZE
    return size


def _deunsync(data):
    """Undo unsynchronisation: every $FF $00 was an $FF before."""
    return data.replace(b'\xff\x00', b'\xff')


# ---------------------------------------------------------------------------
# Frames
# ---------------------------------------------------------------------------

def _plausible_frame_start(buf, at):
    """True if a frame (or padding, or the end of `buf`) can begin at `at`."""
    if at == len(buf):
        return True
    if at > len(buf):
        return False
    head = buf[at:at + 4]
    return bool(_FRAME_ID.match(head)) or not head.strip(b'\x00')


def _v24_frame_size(buf, pos, raw):
    """Size of a v2.4 frame. The spec says synchsafe; some writers (older
    iTunes among them) store a plain integer, so where the two differ, take
    the one after which the next frame starts."""
    plain = int.from_bytes(raw, 'big')
    if any(b & 0x80 for b in raw):
        return plain  # can't be synchsafe
    safe = _synchsafe(raw)
    if safe == plain or _plausible_frame_start(buf, pos + HEADER_SIZE + safe):
        return safe
    if _plausible_frame_start(buf, pos + HEADER_SIZE + plain):
        return plain
    return safe


def _frame_payload(data, fmt, major, unsync_all):
    """The frame's content after its format flags are dealt with, or None if
    it can't be read (encrypted, broken compression)."""
    if major == 3:
        compressed, encrypted = fmt & 0x80, fmt & 0x40
        # Added after the frame header, in flag order: decompressed size (4),
        # encryption method (1), group id (1).
        skip = (4 if compressed else 0) + (1 if encrypted else 0) + (1 if fmt & 0x20 else 0)
        unsync = False  # v2.3 unsynchronises the whole tag, undone already
    else:
        compressed, encrypted = fmt & 0x08, fmt & 0x04
        # Group id (1), encryption method (1), data length indicator (4).
        skip = (1 if fmt & 0x40 else 0) + (1 if encrypted else 0) + (4 if fmt & 0x01 else 0)
        unsync = unsync_all or fmt & 0x02
    if encrypted:
        return None
    data = data[skip:]
    if unsync:
        data = _deunsync(data)
    if compressed:
        try:
            data = zlib.decompressobj().decompress(data, MAX_FRAME_BYTES)
        except zlib.error:
            return None
    return data


def _frames(buf, pos, major, unsync_all=False):
    """(frame id, content) for each complete frame in buf[pos:]."""
    while pos + HEADER_SIZE <= len(buf):
        frame_id = buf[pos:pos + 4]
        if not _FRAME_ID.match(frame_id):
            return  # padding, or something we can't make sense of
        raw_size = buf[pos + 4:pos + 8]
        size = _v24_frame_size(buf, pos, raw_size) if major >= 4 else int.from_bytes(raw_size, 'big')
        fmt = buf[pos + 9]
        start = pos + HEADER_SIZE
        if start + size > len(buf):
            return  # cut off (the caller read only part of a large tag)
        payload = _frame_payload(buf[start:start + size], fmt, major, unsync_all)
        if payload is not None:
            yield frame_id.decode('ascii'), payload
        pos = start + size


# ---------------------------------------------------------------------------
# Frame contents
# ---------------------------------------------------------------------------

_CODECS = {0: 'latin-1', 1: 'utf-16', 2: 'utf-16-be', 3: 'utf-8'}


def _decode(raw, encoding):
    codec = _CODECS.get(encoding)
    if codec is None:
        return ''
    if encoding in (1, 2) and len(raw) % 2:
        raw = raw[:-1]
    return raw.decode(codec, errors='replace')


def _text_frame(data):
    """The (first) string of a T*** frame."""
    if not data:
        return ''
    text = _decode(data[1:], data[0])
    # v2.4 allows several strings, separated by NUL; the first is the title.
    return text.split('\x00')[0].replace('', '').strip()


def _wxxx_frame(data):
    """The URL of a WXXX frame: encoding, description, NUL, URL (latin-1)."""
    if not data:
        return ''
    encoding, rest = data[0], data[1:]
    if encoding in (1, 2):  # the description ends with a two-byte NUL
        i = 0
        while i + 1 < len(rest) and not (rest[i] == 0 and rest[i + 1] == 0):
            i += 2
        url = rest[i + 2:]
    else:
        i = rest.find(b'\x00')
        url = rest[i + 1:] if i >= 0 else b''
    return url.split(b'\x00')[0].decode('latin-1').strip()


def _url_frame(data):
    """The URL of any other W*** frame, which is nothing but the URL."""
    return data.split(b'\x00')[0].decode('latin-1').strip()


def _chap(data, major):
    """(element id, start ms, end ms, title, url) of a CHAP frame, or None."""
    nul = data.find(b'\x00')
    if nul < 0 or len(data) < nul + 17:
        return None
    element_id = data[:nul]
    start_ms, end_ms, _start_offset, _end_offset = struct.unpack('>IIII', data[nul + 1:nul + 17])
    title = url = ''
    # Embedded frames, in the same format as the tag's own. Unsynchronisation
    # of the whole tag has already been undone on `data`.
    for frame_id, payload in _frames(data, nul + 17, major):
        if frame_id == 'TIT2' and not title:
            title = _text_frame(payload)
        elif frame_id == 'WXXX' and not url:
            url = _wxxx_frame(payload)
        elif frame_id.startswith('W') and not url:
            url = _url_frame(payload)
    return element_id, start_ms, end_ms, title, url


def _ctoc(data):
    """(element id, top level?, child element ids) of a CTOC frame, or None."""
    nul = data.find(b'\x00')
    if nul < 0 or len(data) < nul + 3:
        return None
    element_id, flags, count = data[:nul], data[nul + 1], data[nul + 2]
    children, pos = [], nul + 3
    for _ in range(count):
        end = data.find(b'\x00', pos)
        if end < 0:
            break
        children.append(data[pos:end])
        pos = end + 1
    return element_id, bool(flags & 0x02), children


def _toc_order(chaps, tocs):
    """Element ids of the chapters the table of contents lists, in its order
    (top-level CTOCs first, nested ones walked depth first); [] if there's no
    usable table of contents."""
    if not tocs:
        return []
    nested = {child for _, _, children in tocs for child in children}
    roots = [t for t in tocs if t[1]] or [t for t in tocs if t[0] not in nested]
    by_id = {t[0]: t for t in tocs}
    order, seen = [], set()

    def walk(toc):
        if toc[0] in seen:
            return  # a loop
        seen.add(toc[0])
        for child in toc[2]:
            if child in chaps:
                if child not in order:
                    order.append(child)
            elif child in by_id:
                walk(by_id[child])

    for root in roots:
        walk(root)
    return order


def chapters(data):
    """[(start ms, end ms, title, url)] of the chapters in the ID3v2 tag at
    the start of `data`, in table-of-contents order where there is one, else
    in tag order.
    [] if there is no tag or it has no chapters; a tag cut short yields the
    chapters in the part that's there."""
    size = tag_size(data)
    if size is None:
        return []
    major, flags = data[3], data[5]
    if major not in (3, 4):
        return []  # v2.2 has no chapter frames
    body = data[HEADER_SIZE:min(len(data), size)]
    unsync = bool(flags & _UNSYNC)
    if major == 3 and unsync:
        body = _deunsync(body)
    pos = 0
    if flags & _EXTENDED and len(body) >= 4:
        if major == 3:  # size of what follows the size field
            pos = 4 + int.from_bytes(body[:4], 'big')
        else:  # synchsafe size of the whole extended header
            pos = _synchsafe(body[:4])

    chaps, order, tocs = {}, [], []
    for frame_id, payload in _frames(body, pos, major, unsync_all=(major == 4 and unsync)):
        if frame_id == 'CHAP':
            chap = _chap(payload, major)
            if chap and chap[0] not in chaps:
                chaps[chap[0]] = chap
                order.append(chap[0])
        elif frame_id == 'CTOC':
            toc = _ctoc(payload)
            if toc:
                tocs.append(toc)

    # The table of contents' order first; chapters it doesn't list are kept
    # (after it) rather than dropped, since writers don't always list all.
    toc = _toc_order(chaps, tocs)
    listed = set(toc)
    order = toc + [e for e in order if e not in listed]
    return [chaps[e][1:] for e in order]