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

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

HOST = "192.168.1.100"
PORT = 21
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 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):
    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")

    # ---- Upload ----
    print("[UPLOAD] Connecting...")
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(60.0)
    sock.connect((host, port))
    recv_line(sock)  # 220
    send_cmd(sock, "USER anonymous")
    send_cmd(sock, "PASS test@test.com")
    send_cmd(sock, "TYPE I")

    # PASV
    resp = send_cmd(sock, "PASV")
    if not resp:
        print("  PASV failed"); sock.close(); return
    pasv_ip, pasv_port = parse_pasv(resp)
    if pasv_port is None:
        print(f"  PASV parse failed: {resp.decode()}"); sock.close(); return

    # Connect data channel
    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}"); sock.close(); return

    # STOR - measure from STOR send to 226 response
    t0 = time.time()
    send_cmd(sock, f"STOR {test_name}")

    # Send all data at once (sendall handles partial sends)
    data.sendall(test_data)
    data.shutdown(socket.SHUT_WR)
    data.close()

    # Wait for response from server (data received + written to disk)
    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

    # Read STOR response
    resp = recv_line(sock)
    stor_code = resp[:3].decode() if resp else "???"

    print(f"  Sent: {sent:,} / {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'}")

    if stor_code != "226":
        print(f"  UPLOAD FAILED (code={stor_code}), skipping download")
        send_cmd(sock, "QUIT")
        sock.close()
        return

    # ---- Download ----
    time.sleep(3.0)

    resp = send_cmd(sock, "PASV")
    if not resp:
        print("  PASV failed"); sock.close(); return
    pasv_ip, pasv_port = parse_pasv(resp)
    if pasv_port is None:
        print(f"  PASV parse failed"); sock.close(); return

    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}"); sock.close(); return

    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

    resp = recv_line(sock)
    retr_code = resp[:3].decode() if resp else "???"

    print(f"\n[DOWNLOAD]")
    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'}")

    # Compare
    print(f"\n[VERIFY]")
    if len(dl_data) == total_bytes and dl_data == test_data:
        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

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

    print(f"\n=== Summary ===")
    print(f"  File size: {total_bytes:,} bytes ({size_kb} KB)")
    print(f"  Upload:    {upload_speed:.1f} KB/s ({upload_time:.1f}s)")
    print(f"  Download:  {download_speed:.1f} KB/s ({download_time:.1f}s)")

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})")
    args = parser.parse_args()
    test_speed(args.host, args.port, args.size)

if __name__ == "__main__":
    main()