# -*- coding: utf-8 -*-
# RETR-side (device-as-sender) abort probe — validates the 2026-09-24 send
# fast-fail fix (commit 13fe184): client starts a RETR, reads PART of the data,
# then aborts the data connection (RST or FIN) while the device is still
# sending. The device's in-flight net_send must now fail immediately with
# CONNRESET (device log: "short send ... (errno=-8)") instead of stalling up
# to NET_SEND_TIMEOUT_MS = 1s and reporting TIMEDOUT (errno=-9, old behavior).
#
# The STOR-side abort probe (ftp_abort_probe.py) cannot validate this path:
# during STOR the device is the receiver (no in-flight device send exists).
#
# Usage: ftp_retr_abort_probe.py [--ip ..] [--port 21] [--rounds 10]
#        [--file-size 262144] [--read 16384] [--gap 2] [--graceful]
# Output: TS <epoch> / ABORT <n> read=<b> ctrl=<reply> alive=<yes|no> followup=<OK|FAIL:x>
#         SUMMARY rounds=<n> pass=<n> fail=<n>
# Exit 0 all follow-ups pass, 1 otherwise.
import argparse, hashlib, io, random, re, socket, struct, sys, time
from ftplib import FTP, all_errors

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

def rd(sock):
    try:
        return sock.recv(4096).decode("latin1").strip()
    except Exception:
        return ""

def raw_login(ip, port, timeout=10.0):
    s = socket.create_connection((ip, port), timeout=timeout)
    rd(s)                                   # 220 banner
    s.sendall(b"USER anonymous\r\n"); rd(s)
    s.sendall(b"TYPE I\r\n"); rd(s)
    return s

def pasv_data_socket(ctl):
    ctl.sendall(b"PASV\r\n")
    r = rd(ctl)
    m = re.search(r"\((\d+),(\d+),(\d+),(\d+),(\d+),(\d+)\)", r)
    if not m:
        raise RuntimeError("PASV failed: " + r)
    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=10.0)

def followup_ok(ip, port, size=4096, tries=3):
    """normal upload + download + sha256 -> True/False.
    Retries the initial connect a few times: the post-abort relisten window
    (close->reopen chip settle, see bench 20260924 finding F1) may refuse a
    SYN for up to ~5s; a refused follow-up connect is an F1-class failure,
    not a send fast-fail failure — retrying decouples the two."""
    data = bytes(random.getrandbits(8) for _ in range(size))
    want = hashlib.sha256(data).hexdigest()
    for t in range(tries):
        try:
            f = FTP(); f.connect(ip, port, timeout=15.0); f.login("anonymous", "x")
            f.storbinary("STOR /rb_ok.bin", io.BytesIO(data))
            buf = bytearray(); f.retrbinary("RETR /rb_ok.bin", buf.extend)
            try:
                f.delete("/rb_ok.bin")
            except all_errors:
                pass
            f.quit()
            return hashlib.sha256(bytes(buf)).hexdigest() == want
        except all_errors:
            time.sleep(1.0)
    return False

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("--rounds", type=int, default=10)
    ap.add_argument("--file-size", type=int, default=262144,
                    help="size of the file uploaded before the RETR aborts")
    ap.add_argument("--read", type=int, default=16384,
                    help="bytes the client reads before aborting")
    ap.add_argument("--gap", type=float, default=2.0)
    ap.add_argument("--graceful", action="store_true",
                    help="close the data connection with FIN instead of RST")
    a = ap.parse_args()

    # Upload the RETR source file once (8.3 name, FatFs constraint).
    file_data = bytes(random.getrandbits(8) for _ in range(a.file_size))
    f = FTP(); f.connect(a.ip, a.port, timeout=15.0); f.login("anonymous", "setup")
    f.storbinary("STOR /raba.bin", io.BytesIO(file_data))
    f.quit()

    npass = nfail = 0
    try:
        for n in range(1, a.rounds + 1):
            log("TS %d" % int(time.time()))
            ctl_reply = "NA"
            alive = "no"
            got = 0
            try:
                ctl = raw_login(a.ip, a.port)
                ds = pasv_data_socket(ctl)
                ctl.sendall(b"RETR /raba.bin\r\n")
                ctl_reply = rd(ctl)                       # 150
                # read PART of the transfer, then abort while device is sending
                try:
                    got = len(ds.recv(a.read))
                except Exception:
                    pass
                if a.graceful:
                    ds.shutdown(socket.SHUT_RDWR)      # graceful FIN
                else:
                    ds.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER,
                                   struct.pack("ii", 1, 0))   # RST
                ds.close()
                time.sleep(0.5)
                rd(ctl)                                  # 450/226/221 reply
                alive = "yes"
                try:
                    ctl.sendall(b"NOOP\r\n"); rd(ctl)
                except Exception:
                    alive = "no"
            except Exception as e:
                ctl_reply = "EXC:%s" % type(e).__name__
            finally:
                try:
                    ctl.close()
                except Exception:
                    pass
            ok = followup_ok(a.ip, a.port)
            npass += 1 if ok else 0
            nfail += 0 if ok else 1
            log("ABORT %d read=%d ctrl=%s alive=%s followup=%s"
                % (n, got, ctl_reply, alive, "OK" if ok else "FAIL"))
            time.sleep(a.gap)
        f = FTP(); f.connect(a.ip, a.port, timeout=15.0); f.login("anonymous", "cleanup")
        try:
            f.delete("/raba.bin")
        except all_errors:
            pass
        f.quit()
    except Exception as e:
        log("FATAL: %s" % e)

    log("SUMMARY rounds=%d pass=%d fail=%d" % (a.rounds, npass, nfail))
    sys.exit(0 if nfail == 0 else 1)

main()