STM32F4-Base/test/host/ftp_speed_test.py

#!/usr/bin/env python3
"""
FTP 速度测试 — 上传/下载大文件，测量传输速度
用法：python ftp_speed_test.py --host <设备IP> [--size 512] [--port 21]
      [--user admin] [--password admin]
"""

import socket
import sys
import time
import re
import argparse
import os

import bench_defaults  # 台架连接默认值单一真源

HOST = bench_defaults.DEVICE_IP
PORT = 21
# 设备侧 FTP 鉴权走 /config.ini 的 [ftp]（user / pwd_hash），出厂默认 admin/admin；
# **没有匿名通道**（USER anonymous 一律回 530）。
USER = "admin"
PASSWORD = "admin"
FILE_SIZE_KB = 512  # 默认 512KB 测试文件

def log(msg, *args):
    print(f"  {msg % args if args else msg}")

def recv_line(sock, timeout=60.0):
    data = b""
    sock.settimeout(timeout)
    while True:
        try:
            ch = sock.recv(1)
        except socket.timeout:
            return None
        if not ch:
            return data if data else None
        data += ch
        if data.endswith(b"\r\n"):
            return data

def send_cmd(sock, cmd):
    sock.sendall((cmd + "\r\n").encode())
    time.sleep(0.1)
    return recv_line(sock)

def connect_with_retry(host, port, timeout=15.0, attempts=6, delay=3.0):
    """连接控制通道（带重试）。

    CH395F 单客户端复用：上一会话未 QUIT 就断开时，服务端可能仍在处理旧连接，
    新连接会被 RST（WinError 10061）；实测约 30s 后服务端丢弃旧会话并重新 listen。
    注意不要用裸 TCP 探测端口（探测会占掉唯一会话）。
    """
    last_err = None
    for i in range(1, attempts + 1):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(timeout)
        try:
            s.connect((host, port))
            return s, None
        except Exception as e:
            last_err = e
            try:
                s.close()
            except Exception:
                pass
            if i < attempts:
                print(f"  connect failed ({e}); retry {i}/{attempts - 1} in {delay:.0f}s "
                      f"(single-client: stale session blocks new connections)")
                time.sleep(delay)
    return None, last_err


def parse_pasv(resp):
    m = re.search(r"(\d+),(\d+),(\d+),(\d+),(\d+),(\d+)", resp.decode())
    if m:
        a, b, c, d, e, f = [int(x) for x in m.groups()]
        return f"{a}.{b}.{c}.{d}", e * 256 + f
    return None, None

def test_speed(host, port, size_kb, user, password):
    total_bytes = size_kb * 1024
    test_data = os.urandom(total_bytes)
    test_name = "SPEEDTST.DAT"  # 8.3 format

    print(f"\n=== FTP Speed Test: {host}:{port}, {size_kb}KB file ===")
    print(f"  Test data: {total_bytes:,} bytes generated\n")

    checks = []

    def check(label, ok, detail=""):
        checks.append((label, ok, detail))
        mark = "PASS" if ok else "FAIL"
        print(f"    [{mark}] {label}{(' - ' + detail) if detail else ''}")

    def finish(sock):
        passed = all(ok for _, ok, _ in checks)
        print("\n" + "=" * 40)
        print("FTP Speed Test Summary:")
        for label, ok, _ in checks:
            print(f"  [{'PASS' if ok else 'FAIL'}] {label}")
        print("=" * 40)
        if passed:
            print(">>> TEST PASSED <<<")
        else:
            print(">>> TEST FAILED <<<")
        print("=== Test Complete ===")
        if sock:
            try:
                sock.close()
            except Exception:
                pass
        return passed

    # ---- Connect + Login ----
    print("[Connect] Connecting...")
    sock, err = connect_with_retry(host, port, timeout=60.0)
    if sock is None:
        print(f"  CONNECT FAILED: {err}")
        check("Connect", False, str(err))
        return finish(None)
    recv_line(sock)  # 220
    # 需要口令的服务器在 USER 之后回 331，直接放行的才回 230——两种都要接受
    resp = send_cmd(sock, "USER " + user)
    if not (resp and (resp.startswith(b"230") or resp.startswith(b"331"))):
        check("Login USER (230/331)", False,
              resp.decode(errors='replace').rstrip() if resp else "no resp")
        print("  提示：lftpd 按 /config.ini 的 [ftp] 段校验（user / pwd_hash），出厂默认 admin/admin；"
              "设备没有匿名通道。口令被改过就用 --user/--password 传；连续 5 次失败会锁定 60 秒。")
        return finish(sock)
    resp = send_cmd(sock, "PASS " + password)
    if not (resp and resp.startswith(b"230")):
        check("Login PASS 230", False,
              resp.decode(errors='replace').rstrip() if resp else "no resp")
        print("  提示：lftpd 按 /config.ini 的 [ftp] 段校验（user / pwd_hash），出厂默认 admin/admin；"
              "设备没有匿名通道。口令被改过就用 --user/--password 传；连续 5 次失败会锁定 60 秒。")
        return finish(sock)
    check("Login (USER/PASS 230)", True)
    send_cmd(sock, "TYPE I")

    # ---- Upload ----
    print("[UPLOAD] Connecting...")
    resp = send_cmd(sock, "PASV")
    if not resp:
        check("PASV", False, "no response")
        return finish(sock)
    pasv_ip, pasv_port = parse_pasv(resp)
    if pasv_port is None:
        check("PASV parse", False, resp.decode())
        return finish(sock)
    check("PASV parse", True, f"{pasv_ip}:{pasv_port}")

    data = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    data.settimeout(60.0)
    data.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)  # 关 Nagle，防 FIN 跑在数据前面
    try:
        data.connect((pasv_ip, pasv_port))
    except Exception as e:
        print(f"  Data connect failed: {e}")
        check("PASV data connect", False, str(e))
        return finish(sock)
    check("PASV data connect", True)

    t0 = time.time()
    send_cmd(sock, f"STOR {test_name}")
    data.sendall(test_data)
    data.shutdown(socket.SHUT_WR)
    data.close()

    resp = recv_line(sock, timeout=120.0)
    t1 = time.time()
    upload_time = t1 - t0
    upload_speed = total_bytes / upload_time / 1024 if upload_time > 0 else 0
    stor_code = resp[:3].decode() if resp else "???"

    print(f"  Sent: {total_bytes:,} / {total_bytes:,} bytes")
    print(f"  Time: {upload_time:.1f}s")
    print(f"  Speed: {upload_speed:.1f} KB/s")
    print(f"  Response: {resp.decode().rstrip() if resp else 'NONE'}")

    upload_ok = (stor_code == "226")
    check("STOR 226", upload_ok, f"{stor_code} up={upload_speed:.1f}KB/s")

    if not upload_ok:
        send_cmd(sock, "QUIT")
        return finish(sock)

    # ---- Download ----
    time.sleep(3.0)
    print("[DOWNLOAD] Connecting...")
    resp = send_cmd(sock, "PASV")
    if not resp:
        check("RETR PASV", False, "no response")
        return finish(sock)
    pasv_ip, pasv_port = parse_pasv(resp)
    if pasv_port is None:
        check("RETR PASV parse", False, resp.decode())
        return finish(sock)
    data = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    data.settimeout(60.0)
    try:
        data.connect((pasv_ip, pasv_port))
    except Exception as e:
        print(f"  Data connect failed: {e}")
        check("RETR data connect", False, str(e))
        return finish(sock)
    check("RETR data connect", True)

    t0 = time.time()
    send_cmd(sock, f"RETR {test_name}")
    dl_data = b""
    while True:
        try:
            chunk = data.recv(8192)
            if not chunk:
                break
            dl_data += chunk
        except socket.timeout:
            break
    data.close()
    resp = recv_line(sock, timeout=120.0)
    t1 = time.time()
    download_time = t1 - t0
    download_speed = len(dl_data) / download_time / 1024 if download_time > 0 else 0

    print(f"  Received: {len(dl_data):,} / {total_bytes:,} bytes")
    print(f"  Time: {download_time:.1f}s")
    print(f"  Speed: {download_speed:.1f} KB/s")
    print(f"  Response: {resp.decode().rstrip() if resp else 'NONE'}")

    retr_ok = bool(resp) and resp.startswith(b"226")
    check("RETR 226", retr_ok,
          f"{resp[:3].decode() if resp else '???'} dl={download_speed:.1f}KB/s")

    # Compare
    print(f"\n[VERIFY]")
    match = (len(dl_data) == total_bytes and dl_data == test_data)
    if match:
        print(f"  *** MATCH! {total_bytes:,} bytes identical. ***")
    else:
        print(f"  MISMATCH: sent={total_bytes}, got={len(dl_data)}")
        if len(dl_data) != total_bytes:
            print(f"  Size mismatch!")
        else:
            for i in range(total_bytes):
                if dl_data[i] != test_data[i]:
                    print(f"  First diff at byte {i}")
                    break
    check("Upload/Download MATCH", match,
          f"up={total_bytes:,} dl={len(dl_data):,} up={upload_speed:.1f}KB/s dl={download_speed:.1f}KB/s")

    # Cleanup
    send_cmd(sock, f"DELE {test_name}")
    send_cmd(sock, "QUIT")

    return finish(sock)


def main():
    parser = argparse.ArgumentParser(description="FTP Speed Test")
    parser.add_argument("--host", default=HOST)
    parser.add_argument("--port", type=int, default=PORT)
    parser.add_argument("--size", type=int, default=FILE_SIZE_KB,
                        help=f"File size in KB (default {FILE_SIZE_KB})")
    parser.add_argument("--user", default=USER, help=f"FTP user (default: {USER})")
    parser.add_argument("--password", default=PASSWORD,
                        help="FTP password (default: the factory password)")
    args = parser.parse_args()
    success = test_speed(args.host, args.port, args.size, args.user, args.password)
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()