kscheck.py

#!/usr/bin/env python3
"""kscheck.py - compile tableau_hyper.ksy and check it against hyperparse.

An untested .ksy rots. This one had `creator_build` as a u8 at 0x20 long after
README correction 11 retracted that, called the tombstone list "1 byte in every
sample", and had `relation_sample` down as "not observed" - because nothing ever
ran it against a file.

So: compile the spec, parse every `.hyper` in the tree with the generated
parser, and assert it agrees with `hyperparse` field by field. The container
layer is checked positionally; the object types are applied to the decompressed
payloads, which is the only way to reach them (Kaitai has no LZ4).

Needs `kaitai-struct-compiler` on PATH and the `kaitaistruct` runtime:

    brew install kaitai-struct-compiler
    pip install kaitaistruct

    python3 kscheck.py [--keep] [--all-targets]

`--keep` leaves the generated parser behind. `--all-targets` additionally
compiles the spec to every language the compiler supports, which is how you
find out that a construct only works on the target you happened to test.
"""
import glob
import io
import os
import shutil
import subprocess
import sys
import tempfile

import hyperparse as H

HERE = os.path.dirname(os.path.abspath(__file__))
KSY = os.path.join(HERE, "tableau_hyper.ksy")


def compile_spec(outdir):
    exe = shutil.which("kaitai-struct-compiler")
    if not exe:
        sys.exit("kaitai-struct-compiler not on PATH; brew install kaitai-struct-compiler")
    r = subprocess.run([exe, "-t", "python", "--outdir", outdir, KSY],
                       capture_output=True, text=True)
    errs = [l for l in (r.stdout + r.stderr).splitlines() if "error" in l.lower()]
    if r.returncode or errs:
        print(r.stdout + r.stderr)
        sys.exit("the spec does not compile")
    warn = sum(1 for l in (r.stdout + r.stderr).splitlines() if "warning" in l)
    # The remaining warnings are ksy_style_guide naming suggestions that would
    # rename fields away from the engine's own names - `dict_offset` really is
    # an offset the engine calls index.dictOffset, and `num_attributes` counts
    # attributes rather than any one array. Kept for fidelity, deliberately.
    print(f"compiled, no errors ({warn} declined style suggestions)")


# Every target kaitai-struct-compiler 0.11 supports. The spec uses a recursive
# type, parametric types, `_parent` inside value instances and `_io.size`, none
# of which are guaranteed to survive everywhere, so it is worth proving.
TARGETS = ["python", "java", "cpp_stl", "csharp", "javascript", "go", "rust",
           "ruby", "php", "perl", "lua", "nim", "construct", "graphviz", "html"]
# Perl is the one target kaitai-struct-compiler 0.11 has no custom-process
# support for, and the spec needs one to decompress an object. Expected.
UNSUPPORTED = {"perl"}


def compile_all_targets(outdir):
    exe = shutil.which("kaitai-struct-compiler")
    bad = []
    for t in TARGETS:
        r = subprocess.run([exe, "-t", t, "--outdir", os.path.join(outdir, t), KSY],
                           capture_output=True, text=True)
        if r.returncode or any("error" in l.lower()
                               for l in (r.stdout + r.stderr).splitlines()):
            bad.append(t)
    unexpected = [t for t in bad if t not in UNSUPPORTED]
    print(f"targets: {len(TARGETS) - len(bad)}/{len(TARGETS)} compile"
          f" ({sorted(UNSUPPORTED & set(bad))} cannot, as expected)"
          + (f"; UNEXPECTED FAILURES: {unexpected}" if unexpected else ""))
    return unexpected


def check_container(k, hf):
    """Positional fields, against the live root record and the directory."""
    lr, sb = k.live_root, hf.live
    out = {
        "creator": f"{lr.creator_version.major}.{lr.creator_version.minor}."
                   f"{lr.creator_version.build}" == sb.creator_version,
        "min_version": f"{lr.min_version.major}.{lr.min_version.minor}."
                       f"{lr.min_version.build}" == sb.min_version,
        "format": lr.database_format_version == sb.format_version,
        "txn_id": lr.txn_id == sb.txn_id,
        "file_size": lr.file_size == sb.file_size,
        "dir_offset": lr.dir_offset == sb.dir_offset,
        "dir_length": lr.dir_length == sb.dir_length,
        "dir_capacity": lr.dir_capacity == sb.dir_capacity,
        "commit_id": lr.commit_id == sb.commit_id,
        "checksum": lr.checksum == sb.checksum,
    }
    used = [e for e in k.directory.entries if e.state.value == 1]
    out["entry_count"] = len(used) == len(hf.directory.entries)
    out["entries"] = (
        {(e.category, e.index, e.relation, e.block, e.offset, e.size,
          e.compression) for e in hf.directory.entries} ==
        {(e.category.value, e.index, e.relation, e.block, e.pos.offset, e.size,
          e.compression.value) for e in used})
    out["free_list"] = ([(f.offset, f.length) for f in k.directory.free_list]
                        == hf.directory.free_list)
    return out


def _tombstones(node, T, out):
    if node.kind == T.TableauHyper.TombstoneKind.leaf:
        out.extend((r.first, r.last) for r in node.ranges)
    else:
        for c in node.children:
            _tombstones(c, T, out)


def check_objects(T, d, hf, counts, k):
    """The object types, on payloads the SPEC itself reached and decompressed.

    `body` / `body_lz4` do the seeking and the LZ4, so this exercises the custom
    process too - and asserts the bytes match what hyperparse's own decoder
    produces, which is the part that would silently rot if the process were
    wrong.
    """
    from kaitaistruct import KaitaiStream

    def parse(cls, blob):
        return cls(KaitaiStream(io.BytesIO(blob)))

    by_pos = {}
    for ke in k.directory.entries:
        if ke.state.value == 1:
            by_pos[(ke.pos.offset, ke.size)] = (
                ke.body_lz4 if ke.compression.value == 1 else ke.body)

    bad = []
    for e in hf.directory.entries:
        cat = e.category_name
        try:
            blob = H.read_object(d, e)
        except Exception:
            continue
        spec_blob = by_pos.get((e.offset, e.size))
        if spec_blob != blob:
            bad.append(f"{cat}: spec payload differs from hyperparse")
            continue
        try:
            if cat == "Relation_Header":
                o = parse(T.TableauHyper.RelationHeader, blob)
                n = len(H.column_objects(hf, e.relation, 1))
                assert len(o.per_block) == n, f"{len(o.per_block)} u4 for {n} blocks"
                counts["relation_header"] += 1
            elif cat == "Relation_Metadata":
                o = parse(T.TableauHyper.RelationMetadata, blob)
                want, note = H.parse_relation_metadata(blob)
                assert not note, note
                got = []
                if o.present:
                    _tombstones(o.root, T, got)
                assert got == want, f"{len(got)} ranges vs {len(want)}"
                counts["relation_metadata"] += 1
            elif cat == "Relation_DataBlock" and len(blob) >= H.SMA_OFF:
                o = parse(T.TableauHyper.DataBlock, blob)
                ix = H.parse_block_index(blob)
                assert o.tuple_count == ix.tuple_count, "tuple_count"
                assert o.scheme.value == ix.scheme, "scheme"
                assert len(o.sma) == ix.sma_size, "sma"
                assert len(o.dictionary) == ix.dict_size, "dict"
                assert len(o.codes) == ix.data_size, "codes"
                assert len(o.string_heap) == len(ix.string_data(blob)), "heap"
                counts["data_block"] += 1
            elif cat == "Relation_Sample":
                o = parse(T.TableauHyper.RelationSample, blob)
                s = H.parse_relation_sample(blob)
                assert o.num_attributes == len(s.block.groups), "attribute count"
                assert o.version == s.version, "version"
                assert list(o.distinct_estimates) == s.distinct_estimates, "estimates"
                assert (getattr(o, "sample_rows", 0) or 0) == s.sample_rows, "rows"
                assert len(list(getattr(o, "moments", []) or [])) == 3 * len(s.moments), \
                    "moment count"
                for i in range(o.num_attributes):
                    g = o.block.groups[i]
                    sma, dic, dat, heap, scheme = s.block.regions(i)
                    assert g.scheme.value == scheme, f"attr {i} scheme"
                    assert (g.len_sma, g.len_dict, g.len_data, g.len_string_data) \
                        == (len(sma), len(dic), len(dat), len(heap)), f"attr {i} regions"
                counts["relation_sample"] += 1
        except Exception as ex:
            bad.append(f"{cat}: {type(ex).__name__}: {ex}")
    return bad


def main():
    keep = "--keep" in sys.argv
    outdir = os.path.join(HERE, "ksgen") if keep else tempfile.mkdtemp()
    os.makedirs(outdir, exist_ok=True)
    sys.path.insert(0, HERE)          # so the generated parser finds hyperlz4
    try:
        import hyperlz4                                              # noqa: F401
    except ImportError:
        sys.exit("the spec's LZ4 process needs python-lz4: pip install lz4")
    compile_spec(outdir)
    target_fails = compile_all_targets(outdir) if "--all-targets" in sys.argv else []
    sys.path.insert(0, outdir)
    import tableau_hyper as T

    files = sorted(glob.glob(os.path.join(HERE, "**", "*.hyper"), recursive=True))
    counts = dict.fromkeys(
        ("relation_header", "relation_metadata", "data_block", "relation_sample"), 0)
    ok = bad = 0
    for p in files:
        name = os.path.relpath(p, HERE)[:52]
        try:
            k = T.TableauHyper.from_file(p)
            hf = H.parse(p)
        except Exception as ex:
            print(f"  PARSE-FAIL {name}: {type(ex).__name__}: {ex}")
            bad += 1
            continue
        if not hf.directory or not hf.live:
            continue
        fails = [n for n, v in check_container(k, hf).items() if not v]
        fails += check_objects(T, open(p, "rb").read(), hf, counts, k)
        if fails:
            print(f"  MISMATCH {name}: {fails[:4]}")
            bad += 1
        else:
            ok += 1

    print(f"\n{ok} files agree with hyperparse, {bad} disagree")
    print("objects parsed through the spec:")
    for k_, v in counts.items():
        print(f"  {k_:<20} {v:>6}")
    if keep:
        print(f"\ngenerated parser left in {outdir}")
    elif os.path.isdir(outdir):
        shutil.rmtree(outdir, ignore_errors=True)
    return 1 if (bad or target_fails) else 0


if __name__ == "__main__":
    sys.exit(main())