#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
FTP Server 测试脚本 — 连接 STM32F4 上的 lftpd 服务器
用法：python ftp_test.py [--host 192.168.1.100] [--port 21]
"""

import socket
import sys
import time
import argparse

HOST = "192.168.1.100"
PORT = 21
TIMEOUT = 10.0

def recv_line(sock):
    """读取 FTP 响应行（以 \r\n 结尾）"""
    data = b""
    sock.settimeout(TIMEOUT)
    while True:
        try:
            ch = sock.recv(1)
        except socket.timeout:
            print(f"[TIMEOUT] after receiving: {data!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):
    """发送 FTP 命令并读取响应"""
    print(f">>> {cmd}")
    sock.sendall((cmd + "\r\n").encode())
    time.sleep(0.1)
    response = recv_line(sock)
    if response:
        print(f"<<< {response.decode(errors='replace').rstrip()}")
    else:
        print(f"<<< [NO RESPONSE / CONNECTION CLOSED]")
    return response

def parse_pasv(response):
    """解析 PASV 227 响应，返回 (ip, port)"""
    import re
    m = re.search(r"(\d+),(\d+),(\d+),(\d+),(\d+),(\d+)", response.decode())
    if m:
        a, b, c, d, e, f = [int(x) for x in m.groups()]
        ip = f"{a}.{b}.{c}.{d}"
        port = e * 256 + f
        return ip, port
    return None, None

def test_ftp(host, port):
    print(f"=== FTP Test: {host}:{port} ===\n")
    checks = []  # (label, ok, detail)

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

    # 1. Connect
    print("[1] Connecting to control channel...")
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(TIMEOUT)
    try:
        sock.connect((host, port))
        print(f"    Connected to {host}:{port}")
        check("Connect", True)
    except Exception as e:
        print(f"    CONNECT FAILED: {e}")
        check("Connect", False, str(e))
        return False

    # 2. Read welcome
    print("\n[2] Reading welcome banner...")
    resp = recv_line(sock)
    if resp:
        print(f"    {resp.decode(errors='replace').rstrip()}")
        check("Welcome 220", resp.startswith(b"220"),
              resp.decode(errors='replace').rstrip())
    else:
        print("    NO WELCOME BANNER")
        check("Welcome 220", False, "no banner")
        sock.close()
        return False

    # 3. Login
    print("\n[3] Login...")
    resp = send_cmd(sock, "USER anonymous")
    if not resp:
        sock.close()
        check("Login USER", False, "no response")
        return False

    resp = send_cmd(sock, "PASS test@test.com")
    if not resp:
        sock.close()
        check("Login PASS", False, "no response")
        return False
    login_ok = b"230" in resp
    check("Login (USER/PASS 230)", login_ok,
          resp.decode(errors='replace').rstrip())

    # 4. PWD
    print("\n[4] PWD...")
    resp = send_cmd(sock, "PWD")
    pwd_ok = bool(resp) and resp.startswith(b"257")
    check("PWD 257", pwd_ok,
          resp.decode(errors='replace').rstrip() if resp else "no response")

    # 5. PASV mode
    print("\n[5] PASV...")
    resp = send_cmd(sock, "PASV")
    pasv_ip, pasv_port = (None, None)
    if resp:
        pasv_ip, pasv_port = parse_pasv(resp)
    if pasv_port is None and resp and b"227" in resp:
        # 227 响应但解析端口失败
        pasv_port = None
    if pasv_port is None:
        # Try EPSV
        print("\n[5b] Trying EPSV...")
        resp = send_cmd(sock, "EPSV")
        if resp and b"229" in resp:
            import re
            m = re.search(r"\(\|\|\|(\d+)\|\)", resp.decode())
            if m:
                pasv_port = int(m.group(1))
                pasv_ip = host
                print(f"    EPSV: {pasv_ip}:{pasv_port}")
    pasv_ok = pasv_port is not None
    check("PASV/EPSV data port", pasv_ok,
          f"{pasv_ip}:{pasv_port}" if pasv_ok else "no port")

    # 6. LIST (correct PASV flow: connect data first, then send LIST)
    if pasv_ok:
        print(f"    Data channel: {pasv_ip}:{pasv_port}")
        print("\n[6] LIST...")
        data_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        data_sock.settimeout(5.0)
        try:
            data_sock.connect((pasv_ip, pasv_port))
            print(f"    Data connected to {pasv_ip}:{pasv_port}")
            resp = send_cmd(sock, "LIST")
            list_open = bool(resp) and resp.startswith(b"150")
            data = b""
            while True:
                try:
                    chunk = data_sock.recv(1024)
                    if not chunk:
                        break
                    data += chunk
                except socket.timeout:
                    break
            print(f"    LIST result ({len(data)} bytes):")
            print(data.decode(errors='replace'))
            data_sock.close()
            resp = recv_line(sock)
            list_close = bool(resp) and resp.startswith(b"226")
            print(f"    LIST response: "
                  f"{resp.decode(errors='replace').rstrip() if resp else 'none'}")
            list_ok = list_open and list_close
            check("LIST 150->226", list_ok,
                  f"open={list_open} close={list_close} bytes={len(data)}")
        except Exception as e:
            print(f"    LIST FAILED: {e}")
            data_sock.close()
            check("LIST 150->226", False, str(e))
    else:
        print("    Skipping LIST (no data port)")

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

    # Summary
    passed = all(ok for _, ok, _ in checks)
    print("\n" + "=" * 40)
    print("FTP 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 ===")
    return passed

def main():
    parser = argparse.ArgumentParser(description="FTP Server Test")
    parser.add_argument("--host", default=HOST, help=f"FTP server IP (default: {HOST})")
    parser.add_argument("--port", type=int, default=PORT, help=f"FTP server port (default: {PORT})")
    args = parser.parse_args()

    success = test_ftp(args.host, args.port)
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()