#!/usr/bin/env python3
# Reusable bench probe (Trap 23 boundary scan; lives in MDK-ARM/bench_diag/).
# Boundary probe for Trap 23 residual: abort an in-flight STOR, wait G ms, then run a
# fresh STOR(4096) and measure BOTH directions INDEPENDENTLY:
#   ab_stored / fu_stored = device-side byte count, authoritative, parsed from the
#                           device log line  "UPLOAD '<name>' <n> bytes OK"
#   fu_retr               = client-side received length + sha256 (fresh session)
# The abort->reopen gap G is swept so the temporal boundary can be read off directly.
#
# Round:
#   1. abort phase  : login -> PASV -> STOR /aNNNN.bin (8192B burst) -> abort (RST or FIN)
#   2. gap          : sleep G ms   (also measured: abort -> next control/data connect)
#   3. follow-up    : login -> PASV -> STOR /fNNNN.bin (4096B, graceful data close) -> 226
#   4. verify       : fresh session -> RETR /fNNNN.bin -> length + sha256 -> DELE
#   5. classify     : UP_TRUNC (device stored != 4096) / DN_TRUNC (stored ok, client short)
#
# --control drops phase 1 entirely (no preceding abort) = baseline defect rate.
import argparse, hashlib, os, random, re, socket, struct, sys, time

DEVLOG = os.path.join("tools", "serial_tool", "state", "serial.log")


def log(m):
    sys.stdout.write(m + "\n")
    sys.stdout.flush()


class Ctl:
    """Minimal raw FTP control channel (no ftplib) so close semantics/timing are exact."""

    def __init__(self, ip, port, timeout=10.0):
        self.s = socket.create_connection((ip, port), timeout=timeout)
        self.f = self.s.makefile("rb")
        self.banner = self.resp()

    def resp(self):
        line = self.f.readline()
        if not line:
            raise EOFError("control closed by peer")
        txt = line.decode("latin1").rstrip("\r\n")
        code = txt[:3]
        if len(txt) > 3 and txt[3] == "-":
            while True:
                l2 = self.f.readline()
                if not l2:
                    break
                t2 = l2.decode("latin1").rstrip("\r\n")
                txt += "\n" + t2
                if t2.startswith(code + " "):
                    break
        return txt

    def cmd(self, line):
        self.s.sendall((line + "\r\n").encode())
        return self.resp()

    def login(self):
        self.cmd("USER anonymous")
        self.cmd("PASS bench@local")
        self.cmd("TYPE I")

    def pasv(self, timeout=10.0):
        r = self.cmd("PASV")
        m = re.search(r"\((\d+),(\d+),(\d+),(\d+),(\d+),(\d+)\)", r)
        if not m:
            raise RuntimeError("PASV failed: %r" % r[:80])
        host = ".".join(m.group(i) for i in (1, 2, 3, 4))
        port = (int(m.group(5)) << 8) | int(m.group(6))
        return socket.create_connection((host, port), timeout=timeout)

    def close(self):
        """Close BOTH the makefile handle and the socket: leaving the makefile open keeps
        the fd alive, no FIN is sent, and the device holds the stale connection until its
        chip TCP timeout (~2.5s of :21 refusals) -- an artefact, not a device defect."""
        try:
            self.f.close()
        except Exception:
            pass
        try:
            self.s.close()
        except OSError:
            pass

    def abort(self, fin):
        """Tear the control channel down with the same semantic as the trigger."""
        if fin:
            try:
                self.s.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass
        else:
            self.s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
        self.close()

    def quit(self):
        try:
            self.cmd("QUIT")
        except Exception:
            pass
        self.close()


def abort_sock(s, fin, plain=False):
    if plain:
        # pure FIN: plain close, no shutdown(SHUT_RD) -- so the peer's FIN arriving later is
        # NOT answered with an RST (which is what shutdown(SHUT_RDWR) causes on Windows)
        try:
            s.close()
        except OSError:
            pass
        return
    if fin:
        try:
            s.shutdown(socket.SHUT_RDWR)
        except OSError:
            pass
    else:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
    try:
        s.close()
    except OSError:
        pass


def connect_ctl(ip, port, trace=False, tries=4, spacing=0.15):
    """Connect the control channel, tolerating a short post-abort refusal window.

    Returns (Ctl, seconds_since_first_attempt_or_None)."""
    t0 = time.monotonic()
    last = None
    for k in range(tries):
        try:
            return Ctl(ip, port), (time.monotonic() - t0)
        except Exception as e:
            last = e
            if trace:
                log("    [trace] connect attempt %d failed: %s" % (k + 1, type(e).__name__))
            time.sleep(spacing)
    raise last


def read_until_eof(s, deadline_s):
    got = bytearray()
    s.settimeout(deadline_s)
    while True:
        try:
            b = s.recv(65536)
        except socket.timeout:
            break
        except OSError:
            break
        if not b:
            break
        got.extend(b)
    return bytes(got)


def devlog_offset(path):
    try:
        return os.path.getsize(path)
    except OSError:
        return 0


def devlog_new(path, off):
    try:
        with open(path, "rb") as f:
            f.seek(off)
            return f.read().decode("utf-8", "replace")
    except OSError:
        return ""


def parse_uploads(text):
    """name -> stored size, from device log 'UPLOAD '<name>' <n> bytes OK'."""
    out = {}
    for m in re.finditer(r"UPLOAD '([^']*)'\s+(\d+) bytes OK", text):
        out[os.path.basename(m.group(1))] = int(m.group(2))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ip", default="192.168.1.100")
    ap.add_argument("--port", type=int, default=21)
    ap.add_argument("--n", type=int, default=8)
    ap.add_argument("--gap", type=int, default=0, help="abort -> reopen gap in ms")
    ap.add_argument("--size", type=int, default=4096, help="follow-up STOR payload bytes")
    ap.add_argument("--abort-send", type=int, default=8192, help="aborted STOR payload bytes")
    ap.add_argument("--fin", action="store_true", help="graceful FIN abort instead of RST")
    ap.add_argument("--ctl-abort", action="store_true",
                    help="also RST the control channel (aggressive, kills :21 for ~19s)")
    ap.add_argument("--fast-teardown", action="store_true",
                    help="close the control channel ~30ms after the abort instead of 500ms")
    ap.add_argument("--followups", type=int, default=1,
                    help="consecutive STORs after one abort (is the damage one-shot?)")
    ap.add_argument("--plain-close", action="store_true",
                    help="close the data socket with a plain close (pure FIN, no SHUT_RD)")
    ap.add_argument("--control", action="store_true", help="no abort phase (baseline)")
    ap.add_argument("--trace", action="store_true", help="per-phase progress + traceback")
    ap.add_argument("--devlog", default=DEVLOG)
    a = ap.parse_args()

    off = devlog_offset(a.devlog)
    payload = bytes(random.getrandbits(8) for _ in range(a.size))
    want = hashlib.sha256(payload).hexdigest()

    res = {"OK": 0, "UP_TRUNC": 0, "DN_TRUNC": 0, "ERR": 0}
    log("== gap sweep: gap=%dms abort=%s control=%s n=%d size=%d"
        % (a.gap, "FIN" if a.fin else "RST", a.control, a.n, a.size))
    for i in range(1, a.n + 1):
        ab_name = "/a%04d.bin" % i
        fu_name = "/f%04d.bin" % i
        t_abort = None
        dt_ctl = dt_data = -1.0
        try:
            # ---- 1. abort phase ----
            if not a.control:
                c = Ctl(a.ip, a.port)
                if a.trace:
                    log("    [trace] abort-phase control connected: %r" % c.banner[:40])
                c.login()
                d = c.pasv()
                c.cmd("STOR " + ab_name)
                d.sendall(bytes(random.getrandbits(8) for _ in range(a.abort_send)))
                abort_sock(d, a.fin, a.plain_close)
                t_abort = time.monotonic()
                if a.ctl_abort:
                    # aggressive variant: RST the control channel too (kills the :21
                    # listener for ~19s until the chip's own TCP timeout fires -- see
                    # device log 'sock1 TIMEOUT' / 'sock1 relisten quarantine')
                    c.abort(True)
                else:
                    # faithful to ftp_abort_probe/real client: data channel aborted,
                    # control channel closed gracefully. --fast-teardown shrinks the
                    # settle sleep 500ms -> 30ms so the follow-up can land inside the
                    # post-abort window (this is the boundary variable under test).
                    time.sleep(0.03 if a.fast_teardown else 0.5)
                    try:
                        c.resp()
                    except Exception:
                        pass
                    c.quit()
            # ---- 2. gap ----
            if a.gap > 0:
                time.sleep(a.gap / 1000.0)
            # ---- 3./4. follow-up STOR(s) + verify ----
            fu_names = [("/f%04d_%d.bin" % (i, k)) if a.followups > 1 else ("/f%04d.bin" % i)
                        for k in range(a.followups)]
            for k, fu_name in enumerate(fu_names):
                c2, t_accept = connect_ctl(a.ip, a.port, a.trace)
                if a.trace:
                    log("    [trace] follow-up control connected: %r" % c2.banner[:40])
                dt_ctl = (time.monotonic() - t_abort) if t_abort else -1.0
                c2.login()
                d2 = c2.pasv()
                dt_data = (time.monotonic() - t_abort) if t_abort else -1.0
                c2.cmd("STOR " + fu_name)
                d2.sendall(payload)
                try:
                    d2.shutdown(socket.SHUT_WR)
                except OSError:
                    pass
                r = c2.resp()
                d2.close()
                c2.quit()
                # verify: fresh session, RETR
                c3, _ = connect_ctl(a.ip, a.port, a.trace)
                if a.trace:
                    log("    [trace] verify control connected: %r" % c3.banner[:40])
                c3.login()
                d3 = c3.pasv()
                c3.cmd("RETR " + fu_name)
                got = read_until_eof(d3, 15.0)
                d3.close()
                try:
                    c3.cmd("DELE " + fu_name)
                    if not a.control:
                        c3.cmd("DELE " + ab_name)
                except Exception:
                    pass
                c3.quit()
                retr = len(got)
                sha_ok = hashlib.sha256(got).hexdigest() == want
                # stored sizes are resolved after the run from the device log; make a
                # provisional classification from the client side only
                cls = "OK" if (retr == a.size and sha_ok) else "DN_TRUNC"
                log("  %2d.%d resp=%-4s retr=%5d/%d sha=%-4s t_accept=%6.0fms dt_data=%6.0fms  provisional=%s"
                    % (i, k + 1, r[:3], retr, a.size, "OK" if sha_ok else "BAD",
                       (t_accept or 0.0) * 1e3, dt_data * 1e3, cls))
                res[cls] += 1
        except Exception as e:
            res["ERR"] += 1
            log("  %2d EXC %s: %s" % (i, type(e).__name__, str(e)[:70]))
            if a.trace:
                import traceback
                traceback.print_exc()
        time.sleep(0.3)

    # ---- device-log truth: authoritative stored sizes ----
    up = parse_uploads(devlog_new(a.devlog, off))
    log("-- device log: %d UPLOAD records" % len(up))
    up_trunc = 0
    for i in range(1, a.n + 1):
        for k in range(a.followups):
            nm = ("f%04d_%d.bin" % (i, k)) if a.followups > 1 else ("f%04d.bin" % i)
            sz = up.get(nm)
            if sz is None:
                log("  %s : (no UPLOAD record)" % nm)
            elif sz != a.size:
                up_trunc += 1
                log("  %s : stored=%d  <== UPLOAD-SIDE TRUNCATION" % (nm, sz))
    ab = [(k, v) for k, v in up.items() if k.startswith("a")]
    if ab:
        log("-- aborted-phase stored sizes (expected partial): %s"
            % ", ".join("%s=%d" % kv for kv in sorted(ab)))
    log("SUMMARY gap=%dms abort=%s control=%s n=%d client_ok=%d client_dn_trunc=%d "
        "device_up_trunc=%d err=%d"
        % (a.gap, "FIN" if a.fin else "RST", a.control, a.n, res["OK"],
           res["DN_TRUNC"], up_trunc, res["ERR"]))


main()