#!/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 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 ftp_transfer_data(host, pasv_ip, pasv_port, cmd, data_to_send=None, read_data=False):
    """通用数据通道传输：发送命令 + 连数据端口 + 收发数据"""
    print(f"\n  [{cmd}] Opening data channel {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 channel connected.")

        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(TIMEOUT)
        sock.connect((host, PORT))
        recv_line(sock)  # 220
        send_cmd(sock, "USER anonymous")
        send_cmd(sock, "PASS test@test.com")

        # PASV
        resp = send_cmd(sock, "PASV")
        _, data_port = parse_pasv(resp)
        print(f"  PASV port: {data_port}")

        # Connect data FIRST
        data_sock2 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        data_sock2.settimeout(10.0)
        data_sock2.connect((pasv_ip, data_port))

        # Send command
        time.sleep(0.2)
        send_cmd(sock, cmd)

        result = None
        if data_to_send:
            print(f"  Sending {len(data_to_send)} bytes to data channel...")
            total = 0
            while total < len(data_to_send):
                n = data_sock2.send(data_to_send[total:])
                if n <= 0:
                    break
                total += n
            print(f"  Sent {total} bytes.")
            data_sock2.shutdown(socket.SHUT_WR)

        if read_data:
            result = b""
            while True:
                try:
                    chunk = data_sock2.recv(4096)
                    if not chunk:
                        break
                    result += chunk
                except socket.timeout:
                    break
            print(f"  Received {len(result)} bytes from data channel.")

        data_sock2.close()
        resp = recv_line(sock)
        if resp:
            print(f"  Control response: {resp.decode().rstrip()}")

        send_cmd(sock, "QUIT")
        sock.close()
        data_sock.close()
        return result
    except Exception as e:
        print(f"  Data transfer error: {e}")
        try:
            data_sock.close()
        except:
            pass
        return None

def test_upload_download(host, port):
    print(f"=== FTP Upload/Download Test: {host}:{port} ===\n")

    # 1. Connect + Login
    print("[1] Connect + Login")
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(TIMEOUT)
    sock.connect((host, port))
    recv_line(sock)  # 220

    if not send_cmd(sock, "USER anonymous", 230):
        sock.close(); return
    if not send_cmd(sock, "PASS test@test.com", 230):
        sock.close(); return

    # 2. TYPE I (binary mode)
    print("\n[2] Set binary mode")
    send_cmd(sock, "TYPE I")

    # 3. PASV + STOR (upload)
    print("\n[3] Upload: STOR up_test.dat")
    resp = send_cmd(sock, "PASV")
    if not resp:
        sock.close(); return

    pasv_ip, pasv_port = parse_pasv(resp)
    print(f"    PASV: {pasv_ip}:{pasv_port}")

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

    # Send STOR
    send_cmd(sock, "STOR up_test.dat")

    # Send file data
    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()

    # Read STOR response
    resp = recv_line(sock)
    if resp:
        print(f"    Response: {resp.decode().rstrip()}")

    # 4. PASV + LIST (verify)
    print("\n[4] Verify: LIST")
    resp = send_cmd(sock, "PASV")
    if not resp:
        sock.close(); return
    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)
    if resp:
        print(f"    Response: {resp.decode().rstrip()}")

    # 5. PASV + RETR (download)
    time.sleep(1.0)
    print("\n[5] Download: RETR up_test.dat")
    resp = send_cmd(sock, "PASV")
    if not resp:
        sock.close(); return
    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}")
        sock.close(); return

    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)
    if resp:
        print(f"    Response: {resp.decode().rstrip()}")

    # Compare
    print(f"\n[6] Compare")
    print(f"    Uploaded: {len(TEST_DATA)} bytes")
    print(f"    Downloaded: {len(dl_data)} bytes")
    if dl_data == TEST_DATA:
        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}")

    # 7. DELE (delete test file)
    print("\n[7] Cleanup: DELE up_test.dat")
    send_cmd(sock, "DELE up_test.dat")

    # 8. QUIT
    print("\n[8] QUIT")
    send_cmd(sock, "QUIT")
    sock.close()
    print("\n=== Test Complete ===")

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()
    test_upload_download(args.host, args.port)

if __name__ == "__main__":
    main()