#!/usr/bin/env python3
"""Independent check of test_replay's output. See README.md.

For every recorded frame this works out on its own what gn_unwrap_its should have produced -
GeoNetworking and BTP parsed here from EN 302 636-4-1, the TS 103 097 security envelope decoded by
asn1tools from the IEEE 1609.2 ASN.1 modules rather than by hand - and compares that with what the
C code did, record by record. Then it decodes every distinct extracted message with asn1tools and
re-encodes it. Only a message cut at exactly the right byte re-encodes to the same bytes, so this
is what proves that no trailer or signature bytes came along with it.

    py -3.11 check_replay.py replay.tsv capture.pcap [capture.pcap ...]

Needs asn1tools (installed for Python 3.11 on this PC). Message modules come from the repo's
asn1/; the IEEE 1609.2 ones from asn1/ or, failing that, the gitignored vanetza/ checkout.
Exit status 1 if the C code and this disagree about any record.
"""
import collections
import pathlib
import struct
import sys

import asn1tools

ROOT = pathlib.Path(__file__).resolve().parents[3]
PORT_NAMES = {2001: "CAM", 2002: "DENM", 2004: "SPATEM"}
# main.c's RX_FRAME_MAX_LEN: the most of any frame the board hands to gn_unwrap_its. test_replay
# cuts frames to it, so this does too.
RX_FRAME_MAX_LEN = 800


def compile_specs():
    a = ROOT / "asn1"
    uper = {
        2001: asn1tools.compile_files([str(a / "cam_1_4_1.asn"), str(a / "cdd_1_3_1_1.asn")], "uper"),
        2002: asn1tools.compile_files([str(a / "denm_1_3_1.asn"), str(a / "cdd_1_3_1_1.asn")], "uper"),
        2004: asn1tools.compile_files(
            [str(a / n) for n in ("spatem_2_2_1.asn", "mapem_2_2_1.asn", "dsrc_2_2_1.asn", "cdd_2_2_1.asn")],
            "uper"),
    }
    names = ("IEEE1609dot2BaseTypes.asn", "IEEE1609dot2.asn")
    for d in (a, ROOT / "vanetza" / "asn1"):
        if all((d / n).exists() for n in names):
            return uper, asn1tools.compile_files([str(d / n) for n in names], "oer"), d
    return uper, None, None


def frames(path):
    """(record index, 802.11 frame) for every record, numbered the way test_util.c numbers them."""
    d = pathlib.Path(path).read_bytes()
    if len(d) < 24:
        return
    magic = struct.unpack("<I", d[:4])[0]
    e = "<" if magic in (0xA1B2C3D4, 0xA1B23C4D) else ">"
    link = struct.unpack(e + "I", d[20:24])[0]
    if link not in (105, 127):
        return
    off, index = 24, 0
    while off + 16 <= len(d):
        incl = struct.unpack(e + "I", d[off + 8:off + 12])[0]
        if incl > len(d) - off - 16:
            break
        pkt = d[off + 16:off + 16 + incl]
        off += 16 + incl
        if link == 127:
            rl = pkt[2] | pkt[3] << 8 if len(pkt) >= 4 else len(pkt) + 1
            if rl > len(pkt):
                index += 1
                continue
            pkt = pkt[rl:]
        yield index, pkt
        index += 1


def open_envelope(sec, env):
    """(inner GeoNetworking packet, signed) of a TS 103 097 envelope per asn1tools, or None."""
    try:
        m = sec.decode("Ieee1609Dot2Data", env)
    except Exception:
        return None
    if m["protocolVersion"] != 3:
        return None
    kind, content = m["content"]
    if kind == "unsecuredData":
        return content, False
    if kind == "signedData":
        data = content["tbsData"]["payload"].get("data")
        if data and data["protocolVersion"] == 3 and data["content"][0] == "unsecuredData":
            return data["content"][1], True
    return None


def u16(b):
    return int.from_bytes(b, "big")


def s32(b):
    return int.from_bytes(b, "big", signed=True)


def expect(f, sec):
    """What gn_unwrap_its should report for frame f: None, or a dict matching test_replay's row.
    Second value: how many bytes after the message the pre-2026-09-11 code would have forwarded."""
    if len(f) < 24 or (f[0] >> 2) & 3 != 2 or f[1] & 3 == 3:
        return None, None
    o = 24 + (2 if f[0] & 0x80 else 0)
    if f[o:o + 8] != b"\xaa\xaa\x03\x00\x00\x00\x89\x47" or len(f) < o + 12:
        return None, None
    nh = f[o + 8] & 0x0F
    o += 12
    if nh == 2:
        if sec is None:
            return "unchecked", None
        opened = open_envelope(sec, f[o:])
        if opened is None:
            return None, None
        region, signed = opened
    elif nh == 1:
        region, signed = f[o:], False
    else:
        return None, None

    if len(region) < 8 or region[0] >> 4 != 2:
        return None, None
    ht, hst, pl = region[1] >> 4, region[1] & 0x0F, u16(region[4:6])
    if ht == 5 and hst == 0:
        ext, area = 28, None
    elif ht == 4:
        ext = 44
    else:
        return None, None
    if len(region) < 8 + ext + 4:
        return None, None
    if ht == 4:
        a = 8 + 28
        area = (s32(region[a:a + 4]), s32(region[a + 4:a + 8]), u16(region[a + 8:a + 10]))
    port = u16(region[8 + ext:8 + ext + 2])
    if port not in PORT_NAMES or pl <= 4:
        return None, None
    start, end = 8 + ext + 4, 8 + ext + pl
    if signed is not None and nh == 2 and end > len(region):
        return None, None  # the inner packet claims more than its envelope holds
    payload = region[start:min(end, len(region))]
    if not payload:
        return None, None
    old_extra = len(region) - end if nh == 1 else None
    return dict(port=port, signed=signed, truncated=end > len(region), area=area, payload=payload), old_extra


def read_tsv(path):
    rows = {}
    with open(path, encoding="ascii") as fh:
        next(fh)
        for line in fh:
            c = line.rstrip("\n").split("\t")
            key = (c[0], int(c[1]))
            if c[2] == "0":
                rows[key] = None
            else:
                rows[key] = dict(port=int(c[3]), signed=c[4] == "1", truncated=c[5] == "1",
                                 area=(int(c[7]), int(c[8]), int(c[9])) if c[6] == "1" else None,
                                 payload=bytes.fromhex(c[10]))
    return rows


def describe(r):
    if r is None:
        return "rejected"
    return "port %d signed %s truncated %s area %s %d bytes" % (
        r["port"], r["signed"], r["truncated"], r["area"], len(r["payload"]))


def main(argv):
    if len(argv) < 2:
        sys.exit(__doc__)
    got = read_tsv(argv[0])
    uper, sec, sec_dir = compile_specs()
    if sec is None:
        print("WARNING: IEEE 1609.2 modules not found; secured frames are not checked independently")

    stats, extra, disagreements, messages = collections.Counter(), collections.Counter(), [], {}
    for path in argv[1:]:
        for index, f in frames(path):
            key = (path, index)
            stats["records"] += 1
            stats["capped"] += len(f) > RX_FRAME_MAX_LEN
            want, old_extra = expect(f[:RX_FRAME_MAX_LEN], sec)
            if key not in got:
                disagreements.append((key, "no row from test_replay"))
                continue
            have = got.pop(key)
            if want == "unchecked":
                stats["secured, unchecked"] += 1
                continue
            if want != have:
                disagreements.append((key, "C: %s | independent: %s" % (describe(have), describe(want))))
                continue
            if want:
                stats["accepted"] += 1
                stats[PORT_NAMES[want["port"]]] += 1
                stats["signed"] += want["signed"]
                stats["truncated"] += want["truncated"]
                if old_extra is not None:
                    extra[old_extra] += 1
                messages.setdefault((want["port"], want["payload"]), key)
    for key in got:
        disagreements.append((key, "row from test_replay for a record this did not see"))

    print("check_replay: %d records (%d cut to %d bytes), %d accepted (CAM %d, DENM %d, SPATEM %d; "
          "%d signed, %d truncated)"
          % (stats["records"], stats["capped"], RX_FRAME_MAX_LEN, stats["accepted"], stats["CAM"],
             stats["DENM"], stats["SPATEM"], stats["signed"], stats["truncated"]))
    if sec_dir:
        print("  envelopes decoded with asn1tools using %s" % sec_dir.relative_to(ROOT))
    print("  C vs independent parse: %s" % ("agree on every record" if not disagreements
                                             else "%d DISAGREEMENTS" % len(disagreements)))
    for key, why in disagreements[:15]:
        print("    %s record %d: %s" % (pathlib.Path(key[0]).name, key[1], why))

    outcome, failures = collections.Counter(), []
    for (port, payload), key in messages.items():
        name, spec = PORT_NAMES[port], uper[port]
        try:
            decoded = spec.decode(name, payload)
        except Exception as e:
            outcome[(name, "does not decode")] += 1
            failures.append((key, name, len(payload), str(e)[:110]))
            continue
        try:
            encoded = spec.encode(name, decoded)
        except Exception as e:
            # asn1tools decodes an alternative from a later module version (a CHOICE extension)
            # as (None, None) and then cannot encode it back. Says nothing about where the message
            # was cut, so it is reported apart from real mismatches.
            outcome[(name, "decodes; uses a newer extension")] += 1
            failures.append((key, name, len(payload), "cannot re-encode: " + str(e)[:90]))
            continue
        if encoded == payload:
            outcome[(name, "re-encodes byte-identically")] += 1
        else:
            outcome[(name, "re-encodes differently")] += 1
            failures.append((key, name, len(payload), "re-encoded to different bytes"))
    print("  asn1tools on the %d distinct extracted messages:" % len(messages))
    for (name, what), n in sorted(outcome.items()):
        print("    %-6s %-28s %d" % (name, what, n))
    for key, name, n, why in failures[:15]:
        print("    %s record %d, %s %d bytes: %s" % (pathlib.Path(key[0]).name, key[1], name, n, why))
    print("  bytes the pre-2026-09-11 code forwarded after each unsecured message: %s"
          % dict(sorted(extra.items())))
    return 1 if disagreements else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))