#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
pasv_probe.py — decisive probe for the "storm" signature (N34).

Per FTP round prints: device PASV port, PC data-socket SOURCE port, round
result, and timing. Hypothesis under test: the CH395F reuses the same
ephemeral PASV port for the next data listener while the previous data
connection is still in chip-side TIME_WAIT with the same 4-tuple
(PC source port reused too) -> chip reports CONNECT then RSTs ~10ms later
(serial: "sock4 standalone accept" + "DISCONNECTED chip st=0x00/0x00").

Usage:
  python pasv_probe.py --rounds 5 --delay 2          # reproduce (fast rounds)
  python pasv_probe.py --rounds 3 --delay 35         # TIME_WAIT window test
"""
import argparse
import re
import socket
import sys
import time

IP_DEFAULT = "192.168.1.100"


def ftp_stor_round(ip, user, pwd, fname, payload):
    """One login -> PASV -> STOR -> verify transfer. Returns (ok, pasv_port,
    pc_src_port, detail)."""
    c = socket.create_connection((ip, 21), timeout=10)
    c.settimeout(10)
    pc_src = None
    pasv_port = None
    try:
        def readline():
            buf = b""
            while not buf.endswith(b"\n"):
                ch = c.recv(1)
                if not ch:
                    break
                buf += ch
            return buf.decode(errors="replace").strip()

        def cmd(line):
            c.sendall((line + "\r\n").encode())
            return readline()

        greet = readline()
        if not greet.startswith("220"):
            return False, None, None, "no greeting: " + greet
        r = cmd("USER " + user)
        if not r.startswith("331"):
            return False, None, None, "USER: " + r
        r = cmd("PASS " + pwd)
        if not r.startswith("230"):
            return False, None, None, "PASS: " + r
        r = cmd("PASV")
        m = re.search(r"\((\d+),(\d+),(\d+),(\d+),(\d+),(\d+)\)", r)
        if not r.startswith("227") or not m:
            return False, None, None, "PASV: " + r
        g = [int(x) for x in m.groups()]
        pasv_port = g[4] * 256 + g[5]
        d = socket.create_connection((ip, pasv_port), timeout=8)
        pc_src = d.getsockname()[1]
        r = cmd("STOR " + fname)
        if not r.startswith(("125", "150")):
            d.close()
            return False, pasv_port, pc_src, "STOR: " + r
        d.sendall(payload)
        d.close()
        # drain until final 2xx (multi-line or single)
        final = ""
        t0 = time.time()
        while time.time() - t0 < 10:
            r = readline()
            if not r:
                break
            final = r
            if r[:1] in ("2", "4", "5"):
                break
        ok = final.startswith("2")
        cmd("QUIT")
        return ok, pasv_port, pc_src, ("OK" if ok else "final: " + final)
    finally:
        c.close()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ip", default=IP_DEFAULT)
    ap.add_argument("--rounds", type=int, default=5)
    ap.add_argument("--delay", type=float, default=2.0,
                    help="seconds between rounds")
    ap.add_argument("--user", default="admin")
    ap.add_argument("--pwd", default="admin")
    ap.add_argument("--size", type=int, default=1024)
    args = ap.parse_args()

    payload = bytes((i * 7 + 3) & 0xFF for i in range(args.size))
    ports, srcs, oks = [], [], []
    for i in range(args.rounds):
        t = time.strftime("%H:%M:%S")
        ok, port, src, detail = ftp_stor_round(args.ip, args.user, args.pwd,
                                               "PP%d.BIN" % i, payload)
        ports.append(port)
        srcs.append(src)
        oks.append(ok)
        print("[%s] round %d: ok=%-5s pasv_port=%-6s pc_src=%-6s %s"
              % (t, i + 1, ok, port, src, detail))
        if i + 1 < args.rounds:
            time.sleep(args.delay)

    print("\nports :", ports)
    print("pc_src:", srcs)
    print("ok    :", oks)
    same_port = len(set(p for p in ports if p)) == 1
    same_src = len(set(s for s in srcs if s)) == 1
    print("PASV port constant: %s | PC source port constant: %s"
          % (same_port, same_src))
    return 0 if all(oks) else 1


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