pdm/creoparams/lzw.py

"""Decompressione LZW nel formato di compress(1), il vecchio `.Z` Unix.

Creo lo usa per alcuni blocchi interni dei file nativi, fra cui `PictureObj`
che contiene la rappresentazione vettoriale dei disegni.

Rispetto a un LZW generico, compress ha una particolarita' che va riprodotta
fedelmente: i codici sono letti a gruppi di otto e, quando la lunghezza del
codice aumenta o si incontra un CLEAR, i bit residui del gruppo vengono
scartati per riallinearsi. Senza questo dettaglio la decompressione produce
qualche centinaio di byte corretti e poi diverge.
"""

from __future__ import annotations

MAGIC = b"\x1f\x9d"
_CLEAR = 256
_FIRST = 257
INIT_BITS = 9


class NotCompressed(ValueError):
    """I byte non iniziano con la firma di compress."""


def _codes(data: bytes, start: int):
    """Genera (codice, posizione_byte_successiva) applicando le regole di
    allineamento di compress."""
    flags = data[start + 2]
    maxbits = flags & 0x1F
    block_mode = bool(flags & 0x80)

    pos = start + 3
    code_len = INIT_BITS
    next_code = _FIRST if block_mode else 256
    group = []          # codici letti nel gruppo corrente di 8

    while True:
        if not group:
            need = code_len  # 8 codici da code_len bit = code_len byte
            chunk = data[pos : pos + need]
            if len(chunk) < need:
                return
            value = int.from_bytes(chunk, "little")
            mask = (1 << code_len) - 1
            group = [(value >> (i * code_len)) & mask for i in range(8)]
            pos += need

        code = group.pop(0)
        reset = yield code, pos, next_code, code_len
        if reset == "clear":
            group = []
            code_len = INIT_BITS
            next_code = _FIRST
            continue
        if reset is not None:
            next_code = reset
            if next_code >= (1 << code_len) and code_len < maxbits:
                code_len += 1
                group = []  # il cambio di larghezza riallinea il gruppo


def decompress(data: bytes, start: int = 0) -> "tuple[bytes, int]":
    """Decomprime lo stream che inizia a `start`.

    Restituisce (dati, byte_consumati). Si ferma senza sollevare eccezioni al
    primo codice non valido: i blocchi dentro i file Creo non sono terminati
    in modo esplicito, quindi la fine si riconosce cosi'.
    """
    if data[start : start + 2] != MAGIC:
        raise NotCompressed("firma compress (1f 9d) assente")

    block_mode = bool(data[start + 2] & 0x80)
    table: dict[int, bytes] = {i: bytes([i]) for i in range(256)}
    out = bytearray()
    prev: bytes | None = None
    next_code = _FIRST if block_mode else 256
    consumed = start + 3

    gen = _codes(data, start)
    try:
        code, pos, _, _ = next(gen)
        while True:
            if block_mode and code == _CLEAR:
                table = {i: bytes([i]) for i in range(256)}
                next_code = _FIRST
                prev = None
                consumed = pos
                code, pos, _, _ = gen.send("clear")
                continue

            if code in table:
                entry = table[code]
            elif prev is not None and code == next_code:
                entry = prev + prev[:1]
            else:
                break  # fine reale dello stream

            out += entry
            if prev is not None:
                table[next_code] = prev + entry[:1]
                next_code += 1
            prev = entry
            consumed = pos

            code, pos, _, _ = gen.send(next_code)
    except StopIteration:
        pass

    return bytes(out), consumed - start


def find_streams(data: bytes, min_size: int = 256) -> "list[tuple[int, int, bytes]]":
    """Individua e decomprime ogni blocco compress presente nel file.

    Un blocco puo' contenere piu' stream concatenati e la lunghezza consumata
    e' una stima, quindi avanzare a salti farebbe perdere l'inizio di quello
    successivo. Proviamo ogni firma e teniamo gli stream che non ricadono
    dentro uno gia' accettato: cosi' nessun blocco viene saltato e i falsi
    riscontri, che decomprimono pochi byte, cadono da soli.
    """
    candidates = []
    pos = 0
    while True:
        pos = data.find(MAGIC, pos)
        if pos < 0:
            break
        try:
            blob, used = decompress(data, pos)
        except (NotCompressed, IndexError):
            pos += 1
            continue
        if len(blob) >= min_size:
            candidates.append((pos, used, blob))
        pos += 1

    # A parita' di area coperta vince lo stream che produce piu' dati.
    candidates.sort(key=lambda item: (item[0], -len(item[2])))
    found: list = []
    covered_until = -1
    for start, used, blob in candidates:
        if start < covered_until:
            continue
        found.append((start, used, blob))
        covered_until = start + used
    return found