# -*- coding: utf-8 -*-
# Trustworthy abort/robustness probe for the production FTP path (no firmware probes).
#
# Scenario: client starts a STOR upload, sends part of the data, then ABORTS by
# resetting the data connection (SO_LINGER 0 -> RST), which is the common real-world
# pattern (user cancels an upload, client crashes, link drops). Then it checks:
#   1. the control channel still answers (NOOP on a fresh connection),
#   2. a FOLLOW-UP normal upload+download still passes the SHA-256 check,
#   3. the device's relisten actually happened (observable via a successful follow-up).
#
# Usage: ftp_abort_probe.py [--ip ..] [--port 21] [--rounds 10] [--send 8192] [--gap 2]
#        [--graceful] [--check-abort]
# Output: TS <epoch> / ABORT <n> sent=<b> ctrl=<reply> alive=<yes|no> followup=<OK|FAIL:x>
#         abort5xx=<OK|FAIL|->  (only with --check-abort: RST expects 4xx/5xx on the
#         aborted STOR = no false 226 (TRAP-NET-18 hardening); graceful FIN expects 2xx
#         = protocol boundary, still reported as success)
#         SUMMARY rounds=<n> pass=<n> fail=<n>
# Exit 0 all rounds 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
    # N-19: anonymous login is rejected (530) -> use the factory default account
    s.sendall(b"USER admin\r\n"); rd(s)
    s.sendall(b"PASS admin\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):
    """normal upload + download + sha256 -> True/False"""
    data = bytes(random.getrandbits(8) for _ in range(size))
    want = hashlib.sha256(data).hexdigest()
    try:
        f = FTP(); f.connect(ip, port, timeout=15.0); f.login("admin", "admin")
        f.storbinary("STOR /ab_ok.bin", io.BytesIO(data))
        buf = bytearray(); f.retrbinary("RETR /ab_ok.bin", buf.extend)
        try:
            f.delete("/ab_ok.bin")
        except all_errors:
            pass
        f.quit()
        return hashlib.sha256(bytes(buf)).hexdigest() == want
    except all_errors:
        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("--send", type=int, default=8192, help="bytes sent before the RST")
    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")
    ap.add_argument("--check-abort", action="store_true",
                    help="also assert the aborted STOR's control reply: RST -> 4xx/5xx "
                         "(no false 226, TRAP-NET-18 hardening); graceful FIN -> 2xx (boundary)")
    a = ap.parse_args()

    npass = nfail = 0
    for n in range(1, a.rounds + 1):
        log("TS %d" % int(time.time()))
        ctl_reply = "NA"
        alive = "no"
        try:
            ctl = raw_login(a.ip, a.port)
            ds = pasv_data_socket(ctl)
            ctl.sendall(("STOR /ab%04d.bin\r\n" % n).encode())
            ctl_reply = rd(ctl)                       # 150
            chunk = bytes(random.getrandbits(8) for _ in range(a.send))
            ds.sendall(chunk)
            if a.graceful:
                ds.shutdown(socket.SHUT_RDWR)      # graceful FIN (client cancel)
            else:
                # abort: RST the data connection (linger 0)
                ds.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
            ds.close()
            time.sleep(0.5)
            ctl_reply = rd(ctl)                       # 426/550/226
            try:
                ctl.sendall(b"QUIT\r\n"); rd(ctl)
            except Exception:
                pass
            try:
                ctl.close()
            except Exception:
                pass
        except Exception as e:
            ctl_reply = "EXC:%s" % type(e).__name__
        # control-channel liveness on a fresh connection
        try:
            c2 = raw_login(a.ip, a.port)
            c2.sendall(b"NOOP\r\n")
            alive = "yes" if rd(c2).startswith("2") else "no"
            c2.sendall(b"QUIT\r\n"); c2.close()
        except Exception:
            alive = "no"
        fu = followup_ok(a.ip, a.port)
        abort_ok = True
        if a.check_abort:
            code = ctl_reply[0] if (ctl_reply and ctl_reply[0].isdigit()) else "?"
            # RST: device must report 4xx/5xx on the aborted STOR (450); graceful
            # FIN still gets 226 (STOR has no length declaration - protocol boundary)
            abort_ok = (code == "2") if a.graceful else (code in ("4", "5"))
        ok = (alive == "yes" and fu and abort_ok)
        npass += 1 if ok else 0
        nfail += 0 if ok else 1
        log("ABORT %d sent=%d ctrl=%s alive=%s followup=%s abort5xx=%s" %
            (n, a.send, ctl_reply.replace("\r\n", "|")[:40], alive,
             "OK" if fu else "FAIL", "OK" if abort_ok else "FAIL" if a.check_abort else "-"))
        if a.gap > 0 and n < a.rounds:
            time.sleep(a.gap)
    log("SUMMARY rounds=%d pass=%d fail=%d" % (a.rounds, npass, nfail))
    return 0 if nfail == 0 else 1

sys.exit(main())