#!/usr/bin/env python3
"""Attribute captured UTAS requests to labelled UI actions, and diff them.

    utas-filter-diff.py --session <dir> [--baseline NO_FILTER]

The My Squad filter investigation needs to answer "which wire field changed
when I changed exactly one thing in the UI". That is a diff between labelled
groups of requests, so this tool needs both halves:

  raw/utas.ofcap    what the client sent
  labels.txt        MARKER <label> <HH:MM:SS> UTC lines, written when the human
                    said a search was done

Requests are attributed to the label whose marker most recently PRECEDES them.
Anything before the first marker is 'boot'.

Why a separate tool: the observer must stay a dumb, byte-faithful tee. Anything
that interprets traffic belongs outside it, so a mistake here can never affect
what was recorded.
"""
import argparse
import base64
import collections
import hashlib
import json
import os
import re
import sys
from urllib.parse import parse_qsl


def load(session):
    """Rebuild per-connection streams, keeping the wall time of each chunk."""
    path = os.path.join(session, "raw", "utas.ofcap")
    conns = collections.defaultdict(
        lambda: {"c2s": bytearray(), "s2c": bytearray(), "marks": []}
    )
    with open(path) as f:
        for line in f:
            r = json.loads(line)
            if "b64" not in r:
                continue
            c = conns[r["conn"]]
            data = base64.b64decode(r["b64"])
            if r["dir"] == "c2s":
                c["marks"].append((len(c["c2s"]), r["unix"]))
            c[r["dir"]].extend(data)
    return conns


def labels(session):
    out = []
    p = os.path.join(session, "labels.txt")
    if not os.path.exists(p):
        return out
    import datetime
    for line in open(p):
        m = re.match(r"MARKER\s+(\S+)\s+(\d\d):(\d\d):(\d\d)\s+UTC", line.strip())
        if m:
            name, h, mi, s = m.group(1), *map(int, m.groups()[1:])
            out.append((name, h * 3600 + mi * 60 + s))
    return out


def split_requests(buf):
    """(offset, method, target, headers, body) per request in a stream."""
    out, i = [], 0
    while True:
        sep = buf.find(b"\r\n\r\n", i)
        if sep < 0:
            break
        head = bytes(buf[i:sep]).decode("latin1")
        line0 = head.split("\r\n")[0]
        m = re.match(r"(\S+)\s+(\S+)\s+HTTP/", line0)
        if not m:
            break
        hdrs = [h.split(":", 1) for h in head.split("\r\n")[1:] if ":" in h]
        hdrs = [(k.strip(), v.strip()) for k, v in hdrs]
        cl = next((int(v) for k, v in hdrs if k.lower() == "content-length" and v.isdigit()), 0)
        start = sep + 4
        body = bytes(buf[start:start + cl])
        out.append((i, m.group(1), m.group(2), hdrs, body))
        i = start + cl
    return out


def split_responses(buf):
    out, i = [], 0
    while True:
        sep = buf.find(b"\r\n\r\n", i)
        if sep < 0:
            break
        head = bytes(buf[i:sep]).decode("latin1")
        st = re.match(r"HTTP/\d\.\d\s+(\d+)", head)
        cl = re.search(r"(?im)^content-length:\s*(\d+)\s*$", head)
        te = re.search(r"(?im)^transfer-encoding:.*chunked", head)
        start = sep + 4
        if te:
            j, body = start, bytearray()
            while True:
                nl = buf.find(b"\r\n", j)
                if nl < 0:
                    return out
                try:
                    n = int(bytes(buf[j:nl]).split(b";")[0], 16)
                except ValueError:
                    return out
                j = nl + 2
                if n == 0:
                    j = buf.find(b"\r\n", j)
                    j = j + 2 if j >= 0 else len(buf)
                    break
                body.extend(buf[j:j + n])
                j += n + 2
            out.append((int(st.group(1)) if st else None, bytes(body)))
            i = j
        elif cl:
            n = int(cl.group(1))
            out.append((int(st.group(1)) if st else None, bytes(buf[start:start + n])))
            i = start + n
        else:
            out.append((int(st.group(1)) if st else None, b""))
            i = start
    return out


def time_at(marks, off):
    t = marks[0][1] if marks else 0
    for pos, unix in marks:
        if pos > off:
            break
        t = unix
    return t


def item_ids(body):
    """Item ids in a response, if it looks like an item list."""
    try:
        d = json.loads(body)
    except Exception:
        return None
    ids = []

    def walk(o):
        if isinstance(o, dict):
            if "id" in o and isinstance(o["id"], int):
                ids.append(o["id"])
            for v in o.values():
                walk(v)
        elif isinstance(o, list):
            for v in o:
                walk(v)

    walk(d)
    return ids


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--session", required=True)
    ap.add_argument("--baseline", default="NO_FILTER")
    a = ap.parse_args()

    import datetime
    conns = load(a.session)
    marks = labels(a.session)

    rows = []
    for cid in sorted(conns):
        c = conns[cid]
        reqs = split_requests(c["c2s"])
        resps = split_responses(c["s2c"])
        for i, (off, method, target, hdrs, body) in enumerate(reqs):
            unix = time_at(c["marks"], off)
            secs = datetime.datetime.utcfromtimestamp(unix)
            secs = secs.hour * 3600 + secs.minute * 60 + secs.second
            label = "boot"
            for name, at in marks:
                if at <= secs:
                    label = name
            status, rbody = resps[i] if i < len(resps) else (None, b"")
            path, _, query = target.partition("?")
            ids = item_ids(rbody)
            rows.append({
                "label": label,
                "time": datetime.datetime.utcfromtimestamp(unix).strftime("%H:%M:%S"),
                "conn": cid, "method": method, "path": path,
                "query": dict(parse_qsl(query, keep_blank_values=True)) if query else {},
                "query_raw": query,
                "req_body": body.decode("latin1") if len(body) < 2000 else "<%dB>" % len(body),
                "status": status,
                "resp_len": len(rbody),
                "resp_sha": hashlib.sha256(rbody).hexdigest()[:12],
                "item_count": len(ids) if ids is not None else None,
                "item_ids": ids[:40] if ids else None,
            })

    print("=== requests by label ===")
    for r in rows:
        if r["label"] == "boot":
            continue
        print("%-18s %s conn%-3d %-5s %-46s %s %5dB sha=%s items=%s"
              % (r["label"], r["time"], r["conn"], r["method"], r["path"][:46],
                 r["status"], r["resp_len"], r["resp_sha"],
                 r["item_count"] if r["item_count"] is not None else "-"))
        if r["query"]:
            print("%-18s   query: %s" % ("", r["query"]))
        if r["req_body"].strip():
            print("%-18s   body : %s" % ("", r["req_body"][:200]))

    # Field-level diff against the baseline label.
    base = [r for r in rows if r["label"] == a.baseline]
    if not base:
        print("\n(no %s rows yet; skipping diff)" % a.baseline)
        return 0
    print("\n=== query-field diff vs %s ===" % a.baseline)
    bq = base[-1]["query"]
    bpath = base[-1]["path"]
    seen = set()
    for r in rows:
        if r["label"] in ("boot", a.baseline) or r["label"] in seen:
            continue
        seen.add(r["label"])
        added = {k: v for k, v in r["query"].items() if bq.get(k) != v}
        removed = {k: v for k, v in bq.items() if k not in r["query"]}
        note = "" if r["path"] == bpath else "  PATH DIFFERS: %s" % r["path"]
        print("  %-18s +%s  -%s%s" % (r["label"], added or "{}", removed or "{}", note))
    return 0


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