SnailsP2PFile/p2pfile.py

#!/usr/bin/env python3
import os
import sys
import socket
import struct
import json
import hashlib
import argparse
import time
import zipfile
import tarfile
import tempfile
import shutil
import threading

CHUNK_SIZE = 1024 * 1024  # 1MB


def sha256sum(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def pack_folder(folder_path):
    """打包目录，Windows -> zip，Linux -> tar.gz"""
    tmp_dir = tempfile.gettempdir()
    base_name = os.path.basename(os.path.abspath(folder_path))
    if os.name == "nt":
        archive_path = os.path.join(tmp_dir, base_name + ".zip")
        with zipfile.ZipFile(archive_path, "w", zipfile.ZIP_DEFLATED) as zf:
            for root, _, files in os.walk(folder_path):
                for file in files:
                    full = os.path.join(root, file)
                    arcname = os.path.relpath(full, start=folder_path)
                    zf.write(full, arcname)
    else:
        archive_path = os.path.join(tmp_dir, base_name + ".tar.gz")
        with tarfile.open(archive_path, "w:gz") as tf:
            tf.add(folder_path, arcname=base_name)
    return archive_path


def unpack_archive(archive_path, outdir):
    if archive_path.endswith(".zip"):
        with zipfile.ZipFile(archive_path, "r") as zf:
            zf.extractall(outdir)
    elif archive_path.endswith(".tar.gz"):
        with tarfile.open(archive_path, "r:gz") as tf:
            tf.extractall(outdir)

def send_file(conn, filepath, rate_limit=None):
    """发送文件（支持限速+断点续传）"""
    filesize = os.path.getsize(filepath)
    checksum = sha256sum(filepath)

    header = {
        "filename": os.path.basename(filepath),
        "filesize": filesize,
        "sha256": checksum,
    }
    header_bytes = json.dumps(header).encode()
    conn.sendall(struct.pack("!I", len(header_bytes)))
    conn.sendall(header_bytes)

    # 接收对方告知已接收大小（用于断点续传）
    resume_bytes = struct.unpack("!Q", conn.recv(8))[0]
    print(f" 从 {resume_bytes} 字节处开始续传")
    
    sent = resume_bytes
    start_time = time.time()
    last_print = start_time

    with open(filepath, "rb") as f:
        f.seek(resume_bytes)
        while True:
            chunk = f.read(CHUNK_SIZE)
            if not chunk:
                break
            conn.sendall(chunk)
            sent += len(chunk)

            # 限速控制
            if rate_limit:
                elapsed = time.time() - start_time
                if sent / elapsed > rate_limit:
                    time.sleep(0.01)

            # 打印进度每0.5秒更新一次
            now = time.time()
            if now - last_print >= 0.5 or sent == filesize:
                elapsed = now - start_time
                progress = sent / filesize * 100
                avg_speed = sent / elapsed
                remaining = (filesize - sent) / avg_speed if avg_speed > 0 else 0
                print(f"\r {progress:.2f}% | 已发送 {sent/(1024*1024):.2f}/{filesize/(1024*1024):.2f} MB | "
                      f"平均速率 {avg_speed/(1024*1024):.2f} MB/s | 预计剩余 {remaining:.2f}s", end="")
                last_print = now
                sys.stdout.flush()

    total_time = time.time() - start_time
    print(f"\n 发送完成: {filepath} ({filesize} bytes) | 总耗时: {total_time:.2f}s")
    sys.stdout.flush()


def recv_file(conn, outdir):
    """接收文件（支持断点续传）"""
    header_len = struct.unpack("!I", conn.recv(4))[0]
    header = json.loads(conn.recv(header_len).decode())

    filename = header["filename"]
    filesize = header["filesize"]
    checksum = header["sha256"]

    outpath = os.path.join(outdir, filename)
    tmp_path = outpath + ".part"

    received = 0
    if os.path.exists(tmp_path):
        received = os.path.getsize(tmp_path)
        print(f" 检测到未完成文件，已接收 {received} 字节，将继续下载...")

    # 告诉发送端已接收多少
    conn.sendall(struct.pack("!Q", received))

    start_time = time.time()
    last_print = start_time

    with open(tmp_path, "ab") as f:
        while received < filesize:
            chunk = conn.recv(min(CHUNK_SIZE, filesize - received))
            if not chunk:
                break
            f.write(chunk)
            received += len(chunk)

            # 打印进度每0.5秒更新一次
            now = time.time()
            if now - last_print >= 0.5 or received == filesize:
                elapsed = now - start_time
                progress = received / filesize * 100
                avg_speed = received / elapsed
                remaining = (filesize - received) / avg_speed if avg_speed > 0 else 0
                print(f"\r {progress:.2f}% | 已接收 {received/(1024*1024):.2f}/{filesize/(1024*1024):.2f} MB | "
                      f"平均速率 {avg_speed/(1024*1024):.2f} MB/s | 预计剩余 {remaining:.2f}s", end="")
                sys.stdout.flush()
                last_print = now

    total_time = time.time() - start_time
    print()  # 换行
    if received == filesize and sha256sum(tmp_path) == checksum:
        os.rename(tmp_path, outpath)
        print(f"接收完成: {outpath} | 总耗时: {total_time:.2f}s")
        # 自动解压
        #if outpath.endswith(".zip") or outpath.endswith(".tar.gz"):
        #    print("📦 检测到压缩包，自动解压...")
        #    unpack_archive(outpath, outdir)
        #    print(f"📂 已解压到: {outdir}")
    else:
        print("文件不完整或校验失败")
    sys.stdout.flush()



def run_sender(host, port, path, rate, psk):
    if os.path.isdir(path):
        archive = pack_folder(path)
        print(f"目录已打包: {archive}")
        filepath = archive
        cleanup = True
    else:
        filepath = path
        cleanup = False

    s = socket.socket()
    s.connect((host, port))
    s.sendall(psk.encode().ljust(32, b"\0"))  # PSK验证

    rate_limit = None
    if rate:
        unit = rate[-1].upper()
        value = float(rate[:-1])
        if unit == "K":
            rate_limit = value * 1024
        elif unit == "M":
            rate_limit = value * 1024 * 1024

    send_file(s, filepath, rate_limit)
    s.close()

    if cleanup:
        os.remove(filepath)


def handle_client(conn, addr, outdir, psk):
    try:
        recv_psk = conn.recv(32).strip(b"\0").decode()
        if recv_psk != psk:
            print(f"{addr} PSK 验证失败")
            conn.close()
            return

        recv_file(conn, outdir)
    except Exception as e:
        print(f"{addr} 处理错误: {e}")
    finally:
        conn.close()


def run_receiver(port, outdir, psk, once=False):
    s = socket.socket()
    s.bind(("0.0.0.0", port))
    s.listen(5)
    print(f"正在监听 {port}...")
    sys.stdout.flush()

    while True:
        conn, addr = s.accept()
        print(f"收到连接: {addr}")
        # 用线程处理每个客户端
        t = threading.Thread(target=handle_client, args=(conn, addr, outdir, psk))
        t.daemon = True
        t.start()

        if once:
            print("一次传输模式，退出主循环")
            break

    s.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="P2P 文件/目录传输工具（支持续传）")
    parser.add_argument("--mode", choices=["send", "recv"], required=True)
    parser.add_argument("--host", default="127.0.0.1", help="接收端IP (send模式必需，默认127.0.0.1)")
    parser.add_argument("--port", type=int, help="端口号")
    parser.add_argument("--path", help="要发送的文件/目录 (send模式必需)")
    parser.add_argument("--outdir", help="接收保存目录 (recv模式必需)")
    parser.add_argument("--rate", help="限速，例: 500K, 5M")
    parser.add_argument("--psk", required=True, help="预共享密钥")
    parser.add_argument("--once", action="store_true", help="接收一个文件后退出")
    args = parser.parse_args()

    # 设置默认端口
    if args.port is None:
        if args.mode == "send":
            args.port = 7001
        else:
            args.port = 6000

    if args.mode == "send":
        if not args.path:
            parser.error("send 模式需要 --path")
        run_sender(args.host, args.port, args.path, args.rate, args.psk)
    else:
        if not args.outdir:
            parser.error("recv 模式需要 --outdir")
        os.makedirs(args.outdir, exist_ok=True)
        run_receiver(args.port, args.outdir, args.psk, args.once)