wifile.py

"""
WiFile - A simple command-line tool for transferring files over a network.

This module provides functionality to send and receive files between devices
on the same network using TCP sockets. It operates in two modes:
- Server mode: sends a file to connected clients
- Client mode: receives a file from a server
"""
import socket
import argparse
import os
import sys
import time


def format_bytes(bytes_value):
    """Convert bytes to human readable format."""
    for unit in ['B', 'KB', 'MB', 'GB']:
        if bytes_value < 1024.0:
            return f"{bytes_value:.1f} {unit}"
        bytes_value /= 1024.0
    return f"{bytes_value:.1f} TB"


def show_progress(current, total, start_time=None):
    """Display a progress bar for file transfer."""
    if total == 0:
        return

    percent = (current / total) * 100
    bar_length = 50
    filled_length = int(bar_length * current // total)
    progress_bar = '█' * filled_length + '-' * (bar_length - filled_length)

    current_formatted = format_bytes(current)
    total_formatted = format_bytes(total)

    progress_line = f'\r|{progress_bar}| {percent:.1f}% ({current_formatted}/{total_formatted})'

    if start_time:
        elapsed = time.time() - start_time
        if elapsed > 0 and current > 0:
            speed = current / elapsed
            speed_formatted = format_bytes(speed)
            eta_seconds = (total - current) / speed if speed > 0 else 0
            eta_formatted = f"{int(eta_seconds)}s"
            progress_line += f' - {speed_formatted}/s - ETA: {eta_formatted}'

    print(progress_line, end='', flush=True)

    if current >= total:
        print()  # New line when complete


def get_local_ip():
    """Get the local IP address of the machine."""
    try:
        # Create a socket and connect to a remote address to determine local IP
        # This doesn't actually send data, just determines routing
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            s.connect(("8.8.8.8", 80))
            local_ip = s.getsockname()[0]
        return local_ip
    except (socket.error, OSError):
        # Fallback to localhost if we can't determine the actual IP
        return "127.0.0.1"


def start_server(port, filepath):
    """Run the server to send a file to a client."""
    if not os.path.isfile(filepath):
        print(f"Error: File '{filepath}' does not exist.")
        sys.exit(1)

    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind(('0.0.0.0', port))  # Listen on all interfaces
    server_socket.listen(1)

    local_ip = get_local_ip()
    print(f"Server listening on port {port}")
    print(f"Server IP address: {local_ip}")
    print(
        f"Clients can connect using: python wifile.py client --host {local_ip}")
    print("Waiting for connection...")

    conn = None
    try:
        conn, addr = server_socket.accept()
        conn.settimeout(30)  # 30 second timeout
        print(f"Connected by {addr}")

        # Send file name and size first
        filename = os.path.basename(filepath)
        filesize = os.path.getsize(filepath)
        header = f"{filename}:{filesize}\n".encode()
        conn.send(header)

        # Wait for client acknowledgment
        try:
            ack = conn.recv(4)  # Expect "ACK"
            if ack != b"ACK\n":
                print("Client did not acknowledge header properly")
                return
        except socket.timeout:
            print("Timeout waiting for client acknowledgment")
            return

        # Send file content with progress bar
        print(f"Sending '{filename}' ({format_bytes(filesize)})...")
        sent_bytes = 0
        start_time = time.time()

        with open(filepath, 'rb') as f:
            while True:
                data = f.read(1024)  # Read in 1KB chunks
                if not data:
                    break
                try:
                    conn.send(data)
                    sent_bytes += len(data)
                    show_progress(sent_bytes, filesize, start_time)
                except (socket.error, ConnectionResetError, ConnectionAbortedError, BrokenPipeError) as e:
                    print(f"\nConnection lost during transfer: {e}")
                    transfer_msg = (f"Transfer incomplete: {format_bytes(sent_bytes)} "
                                    f"of {format_bytes(filesize)} sent")
                    print(transfer_msg)
                    return

        print(f"File '{filename}' sent successfully.")
    except (socket.error, OSError, IOError) as e:
        if "10054" in str(e) or "forcibly closed" in str(e).lower():
            print(f"\nClient disconnected unexpectedly: {e}")
            disconnect_msg = ("This usually means the client closed the connection "
                              "or network was interrupted.")
            print(disconnect_msg)
        else:
            print(f"Server error: {e}")
    finally:
        if conn is not None:
            conn.close()
        server_socket.close()


def start_client(host, port, output_dir, auto_overwrite=False, auto_rename=False):
    """Run the client to receive a file from the server."""
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.settimeout(30)  # 30 second timeout
    try:
        client_socket.connect((host, port))
        print(f"Connected to server {host}:{port}")

        # Receive file name and size
        header_data = b""
        while True:
            chunk = client_socket.recv(1)
            if not chunk:
                raise ConnectionError(
                    "Connection closed while receiving header")
            header_data += chunk
            if header_data.endswith(b'\n'):
                break

        try:
            header_str = header_data.decode('utf-8').strip()
            filename, filesize = header_str.split(':')
            filesize = int(filesize)
        except (UnicodeDecodeError, ValueError) as e:
            raise ValueError(f"Invalid header format: {e}") from e

        # Send acknowledgment
        client_socket.send(b"ACK\n")

        output_path = os.path.join(output_dir, filename)

        # Handle file conflicts
        if os.path.exists(output_path):
            if auto_overwrite:
                print(f"Overwriting existing file '{filename}'...")
            elif auto_rename:
                base_name, ext = os.path.splitext(filename)
                counter = 1
                while True:
                    new_filename = f"{base_name}_{counter}{ext}"
                    new_output_path = os.path.join(output_dir, new_filename)
                    if not os.path.exists(new_output_path):
                        output_path = new_output_path
                        filename = new_filename
                        print(f"Saving as '{filename}' to avoid conflict...")
                        break
                    counter += 1
            else:
                print(
                    f"Warning: File '{filename}' already exists in '{output_dir}'")
                while True:
                    choice = input(
                        "Choose action: (o)verwrite, (r)ename, (c)ancel: ").lower().strip()
                    if choice in ['o', 'overwrite']:
                        print(f"Overwriting existing file '{filename}'...")
                        break
                    elif choice in ['r', 'rename']:
                        base_name, ext = os.path.splitext(filename)
                        counter = 1
                        while True:
                            new_filename = f"{base_name}_{counter}{ext}"
                            new_output_path = os.path.join(
                                output_dir, new_filename)
                            if not os.path.exists(new_output_path):
                                output_path = new_output_path
                                filename = new_filename
                                print(f"Saving as '{filename}' instead...")
                                break
                            counter += 1
                        break
                    elif choice in ['c', 'cancel']:
                        print("Transfer cancelled by user.")
                        return
                    else:
                        print("Invalid choice. Please enter 'o', 'r', or 'c'.")

        # Receive file content with progress bar
        print(f"Receiving '{filename}' ({format_bytes(filesize)})...")
        received = 0
        start_time = time.time()

        with open(output_path, 'wb') as f:
            while received < filesize:
                try:
                    data = client_socket.recv(1024)
                    if not data:
                        print("\nConnection closed by server. Transfer incomplete.")
                        received_msg = (f"Received: {format_bytes(received)} "
                                        f"of {format_bytes(filesize)}")
                        print(received_msg)
                        break
                    f.write(data)
                    received += len(data)
                    show_progress(received, filesize, start_time)
                except (socket.error, ConnectionResetError, ConnectionAbortedError) as e:
                    print(f"\nConnection lost during transfer: {e}")
                    received_msg = (f"Received: {format_bytes(received)} "
                                    f"of {format_bytes(filesize)}")
                    print(received_msg)
                    break

        if received == filesize:
            print(f"File '{filename}' received and saved to '{output_path}'.")
        else:
            incomplete_msg = (f"Transfer incomplete. File saved as '{output_path}' "
                              "but may be corrupted.")
            print(incomplete_msg)
    except (socket.error, OSError, IOError, ValueError) as e:
        if "10054" in str(e) or "forcibly closed" in str(e).lower():
            print(f"Server disconnected unexpectedly: {e}")
            disconnect_msg = ("This usually means the server closed the connection "
                              "or network was interrupted.")
            print(disconnect_msg)
        else:
            print(f"Client error: {e}")
    finally:
        client_socket.close()


def main():
    """Parse command-line arguments and run the appropriate mode (server or client)."""
    parser = argparse.ArgumentParser(
        description="CLI tool for file transfer over WiFi network")
    parser.add_argument(
        'mode', choices=['server', 'client'], help="Run as server or client")
    parser.add_argument('--port', type=int, default=12345,
                        help="Port to use (default: 12345)")
    parser.add_argument(
        '--file', help="Path to the file to send (server mode)")
    parser.add_argument('--host', help="Server IP address (client mode)")
    parser.add_argument('--output-dir', default=".",
                        help="Directory to save received file (client mode)")
    parser.add_argument('--overwrite', action='store_true',
                        help="Automatically overwrite existing files (client mode)")
    parser.add_argument('--auto-rename', action='store_true',
                        help="Automatically rename if file exists (client mode)")

    args = parser.parse_args()

    if args.mode == 'server':
        if not args.file:
            print("Error: --file is required in server mode.")
            sys.exit(1)
        start_server(args.port, args.file)
    elif args.mode == 'client':
        if not args.host:
            print("Error: --host is required in client mode.")
            sys.exit(1)
        start_client(args.host, args.port, args.output_dir,
                     args.overwrite, args.auto_rename)


if __name__ == "__main__":
    main()