#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
FTP 上传+下载测试脚本
用法：python ftp_up_dl_test.py --host 192.168.1.100
"""

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

HOST = "192.168.1.100"
PORT = 21
TIMEOUT = 15.0

TEST_DATA = b"Hello from Python FTP test!\r\nThis is line 2.\r\nLine 3: " + b"ABCDEFGH" * 50 + b"\r\n"

def recv_line(sock, timeout=TIMEOUT):
    data = b""
    sock.settimeout(timeout)
    while True:
        try:
            ch = sock.recv(1)
        except socket.timeout:
            print(f"  [TIMEOUT] partial: {data[:80]!r}")
            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, expect_code=None):
    print(f"  > {cmd}")
    sock.sendall((cmd + "\r\n").encode())
    time.sleep(0.15)
    response = recv_line(sock)
    if response:
        text = response.decode(errors='replace').rstrip()
        print(f"  < {text}")
        if expect_code and not text.startswith(str(expect_code)):
            print(f"  *** Expected {expect_code}, got: {text}")
            return None
        return response
    print("  < [NO RESPONSE]")
    return None

def connect_with_retry(host, port, timeout=TIMEOUT, 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_upload_download(host, port):
    print(f"=== FTP Upload/Download Test: {host}:{port} ===\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 Upload/Download 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

    # 1. Connect + Login
    print("[1] Connect + Login")
    sock, err = connect_with_retry(host, port)
    if sock is None:
        print(f"    CONNECT FAILED: {err}")
        check("Connect", False, str(err))
        return finish(None)
    recv_line(sock)  # 220

    if not send_cmd(sock, "USER anonymous", 230):
        check("Login USER 230", False)
        return finish(sock)
    if not send_cmd(sock, "PASS test@test.com", 230):
        check("Login PASS 230", False)
        return finish(sock)
    check("Login (USER/PASS 230)", True)

    # 2. TYPE I (binary mode)
    print("\n[2] Set binary mode")
    resp = send_cmd(sock, "TYPE I")
    type_ok = bool(resp) and resp.startswith(b"200")
    check("TYPE I 200", type_ok)

    # 3. PASV + STOR (upload)
    print("\n[3] Upload: STOR up_test.dat")
    resp = send_cmd(sock, "PASV")
    if not resp:
        check("PASV", False, "no response")
        return finish(sock)
    pasv_ip, pasv_port = parse_pasv(resp)
    print(f"    PASV: {pasv_ip}:{pasv_port}")

    data_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    data_sock.settimeout(10.0)
    try:
        data_sock.connect((pasv_ip, pasv_port))
        print(f"    Data connected.")
    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, f"{pasv_ip}:{pasv_port}")

    send_cmd(sock, "STOR up_test.dat")
    total = 0
    while total < len(TEST_DATA):
        try:
            n = data_sock.send(TEST_DATA[total:])
            if n <= 0:
                break
            total += n
        except Exception as e:
            print(f"    Send error: {e}")
            break
    print(f"    Uploaded {total}/{len(TEST_DATA)} bytes.")
    data_sock.shutdown(socket.SHUT_WR)
    data_sock.close()

    resp = recv_line(sock)
    stor_ok = bool(resp) and resp.startswith(b"226") and total == len(TEST_DATA)
    check("STOR 226 + bytes", stor_ok,
          f"uploaded={total}/{len(TEST_DATA)} resp={resp.decode().rstrip() if resp else 'none'}")

    # 4. PASV + LIST (verify)
    print("\n[4] Verify: LIST")
    resp = send_cmd(sock, "PASV")
    if not resp:
        check("LIST PASV", False, "no response")
        return finish(sock)
    pasv_ip, pasv_port = parse_pasv(resp)

    data_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    data_sock.settimeout(10.0)
    data_sock.connect((pasv_ip, pasv_port))

    send_cmd(sock, "LIST")
    data = b""
    while True:
        try:
            chunk = data_sock.recv(4096)
            if not chunk:
                break
            data += chunk
        except socket.timeout:
            break
    print(f"    LIST: {data.decode(errors='replace').rstrip() or '(empty directory)'}")
    data_sock.close()

    resp = recv_line(sock)
    list_ok = bool(data) and b"UP_TEST.DAT" in data.upper()
    check("LIST shows UP_TEST.DAT", list_ok,
          resp.decode().rstrip() if resp else "no resp")

    # 5. PASV + RETR (download)
    time.sleep(1.0)
    print("\n[5] Download: RETR up_test.dat")
    resp = send_cmd(sock, "PASV")
    if not resp:
        check("RETR PASV", False, "no response")
        return finish(sock)
    pasv_ip, pasv_port = parse_pasv(resp)

    data_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    data_sock.settimeout(10.0)
    try:
        data_sock.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, f"{pasv_ip}:{pasv_port}")

    send_cmd(sock, "RETR up_test.dat")
    dl_data = b""
    while True:
        try:
            chunk = data_sock.recv(4096)
            if not chunk:
                break
            dl_data += chunk
        except socket.timeout:
            break
    data_sock.close()
    resp = recv_line(sock)
    retr_ok = bool(resp) and resp.startswith(b"226")
    check("RETR 226", retr_ok, resp.decode().rstrip() if resp else "no resp")

    # Compare
    print(f"\n[6] Compare")
    print(f"    Uploaded: {len(TEST_DATA)} bytes")
    print(f"    Downloaded: {len(dl_data)} bytes")
    match = (dl_data == TEST_DATA)
    if match:
        print(f"    *** MATCH! Upload/Download success. ***")
    else:
        print(f"    *** MISMATCH! ***")
        if dl_data:
            print(f"    First 80 bytes uploaded: {TEST_DATA[:80]!r}")
            print(f"    First 80 bytes downloaded: {dl_data[:80]!r}")
    check("Upload/Download MATCH", match, f"up={len(TEST_DATA)} dl={len(dl_data)}")

    # 7. DELE (delete test file)
    print("\n[7] Cleanup: DELE up_test.dat")
    resp = send_cmd(sock, "DELE up_test.dat")
    dele_ok = bool(resp) and resp.startswith(b"250")
    check("DELE 250", dele_ok, resp.decode().rstrip() if resp else "no resp")

    # 8. QUIT
    print("\n[8] QUIT")
    resp = send_cmd(sock, "QUIT")
    quit_ok = bool(resp) and resp.startswith(b"221")
    check("QUIT 221", quit_ok, resp.decode().rstrip() if resp else "no resp")

    return finish(sock)


def main():
    parser = argparse.ArgumentParser(description="FTP Upload/Download Test")
    parser.add_argument("--host", default=HOST)
    parser.add_argument("--port", type=int, default=PORT)
    args = parser.parse_args()
    success = test_upload_download(args.host, args.port)
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()