#!/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")

    # 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}")
    except Exception as e:
        print(f"    CONNECT FAILED: {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()}")
    else:
        print("    NO WELCOME BANNER")
        sock.close()
        return False

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

    resp = send_cmd(sock, "PASS test@test.com")
    if not resp:
        sock.close()
        return False

    # 4. PWD
    print("\n[4] PWD...")
    resp = send_cmd(sock, "PWD")

    # 5. PASV mode
    print("\n[5] PASV...")
    resp = send_cmd(sock, "PASV")
    if not resp:
        sock.close()
        return False

    pasv_ip, pasv_port = parse_pasv(resp)
    if pasv_port is None:
        print("    FAILED to parse PASV response")
        # 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}")

    if pasv_port is None:
        print("    PASV/EPSV FAILED, trying PORT mode...")
        # Skip data connection tests
    else:
        print(f"    Data channel: {pasv_ip}:{pasv_port}")

        # 6. LIST (correct PASV flow: connect data first, then send LIST)
        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}")
            send_cmd(sock, "LIST")
            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()
        except Exception as e:
            print(f"    LIST FAILED: {e}")
            data_sock.close()

        # Read LIST response
        resp = recv_line(sock)
        if resp:
            print(f"    LIST response: {resp.decode(errors='replace').rstrip()}")

    # 7. QUIT
    print("\n[7] QUIT...")
    send_cmd(sock, "QUIT")

    sock.close()
    print("\n=== Test Complete ===")
    return True

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()