Python send and reveive files via TCP connection.

tcp-file-transfer.py

#!/usr/bin/env python3

import argparse
import logging
import os
import socket
import sys

BUFFER_SIZE = 1200

UPLOAD_DIRECTORY = './uploads'

FLCT_START =    b'<<__START__>>'
FLCT_FILENAME = b'<<__FILENAME__>>'
FLCT_FILESIZE = b'<<__FILESIZE__>>'
FLCT_BYTES =    b'<<__BYTES__>>'
FLCT_FINISHED = b'<<__FINISHED__>>'
FLCT_END =      b'<<__END__>>'
FLCT_OK =       b'<<__OK__>>'


def _read_message(conn) -> bytes:
    buffer = b''

    while True:
        data = conn.recv(BUFFER_SIZE)
        buffer += data

        if buffer.startswith(FLCT_START) and buffer.endswith(FLCT_END):
            break
        
        if buffer == FLCT_OK:
            break

    return buffer


def client_mode(server_addr: str, server_port: int, absolute_path: str):
    server_sock_data = server_addr, server_port
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect(server_sock_data)

    filename = os.path.basename(absolute_path)
    filesize = os.stat(absolute_path).st_size
    print('INFO file {}, size {} bytes.'.format(filename, filesize))

    message = FLCT_START + FLCT_FILENAME + filename.encode() + FLCT_END
    sock.sendall(message)
    data = _read_message(sock)
    if data != FLCT_OK:
        print('ERROR abort.')

    message = FLCT_START + FLCT_FILESIZE + str(filesize).encode() + FLCT_END
    sock.sendall(message)
    data = _read_message(sock)
    if data != FLCT_OK:
        print('ERROR abort.')
    
    sending_bytes = 0
    percentage_step = 10

    fd = open(absolute_path, 'rb')
    while True:
        chunk = fd.read(BUFFER_SIZE)
        if chunk:
            sock.sendall(FLCT_START + FLCT_BYTES + chunk + FLCT_END)
            data = _read_message(sock)
            if data != FLCT_OK:
                print('ERROR abort.')
            
            sending_bytes += len(chunk)
            percentage = 100 * sending_bytes / filesize
            
            if percentage >= percentage_step:
                print('INFO sending {}%'.format(percentage_step), end='\r')
                percentage_step += 10
        
        else:
            break
    
    print('INFO sending 100%')
    
    sock.sendall(FLCT_START + FLCT_FINISHED + FLCT_END)
    data = _read_message(sock)
    if data != FLCT_OK:
        print('ERROR abort.')
    
    print('INFO file {} sended.'.format(filename))
    sock.close()


def server_mode(port: int):
    os.makedirs(UPLOAD_DIRECTORY, exist_ok=True)
    
    sock = socket.socket(socket.AF_INET,  socket.SOCK_STREAM)
    sock.bind(('0.0.0.0', port))
    sock.listen()
    print('start server, listening on port {}.'.format(port))

    while True:
        conn, addr = sock.accept()
        print('connection accept, address {0[0]}:{0[1]}.'.format(addr))

        filename = ""
        filesize = 0
        fd = None
        
        received_bytes = 0
        percentage_step = 10

        try:
            while True:
                data = _read_message(conn)
                conn.sendall(FLCT_OK)
                
                data = data[len(FLCT_START):-len(FLCT_END)]
                
                if data.startswith(FLCT_FILENAME):
                    filename = data[len(FLCT_FILENAME):].decode()
                    print('INFO open file {}.'.format(filename))
                    fd = open(os.path.join(UPLOAD_DIRECTORY, filename), 'wb')
                    continue

                if data.startswith(FLCT_FILESIZE):
                    filesize = int(data[len(FLCT_FILESIZE):].decode())
                    print('INFO file size {} bytes.'.format(filesize))
                    continue

                if data.startswith(FLCT_BYTES):
                    if fd is None:
                        print('ERROR cannot open filename.', file=sys.stderr)
                        break
                    
                    chunk = data[len(FLCT_BYTES):]
                    received_bytes += len(chunk)
                    fd.write(chunk)

                    percentage = 100 * received_bytes / filesize
                    
                    if percentage >= percentage_step:
                        print('INFO received {}%'.format(percentage_step), end='\r')
                        percentage_step += 10
                    
                if data.startswith(FLCT_FINISHED):
                    fd.close()
                    print('INFO file {} saved.'.format(filename))
                    break
            
            print('INFO received 100%')
            
        except Exception as ex:
            print('ERROR {}: {}.'.format(ex.__class__.__name__, ex))

        print('INFO connection close.')
        conn.close()


def main():
    parser = argparse.ArgumentParser(description='Python File Transfer.')
    parser.add_argument('--server', '-s', dest='server', action='store_true', default=False, help='server mode, receive files.')
    parser.add_argument('--addr', '-a', dest='addr', type=str, action='store', default='', help='client mode, ip to send the file.')
    parser.add_argument('--port', '-p', dest='port',  type=int, action='store', default=9999, help='[client|server] mode, port to send or receive the file.')
    parser.add_argument('--file', '-f', dest='file', type=str, action='store', default='', help='client mode, filename.')
    args = parser.parse_args(sys.argv[1:])

    if args.server:
        server_mode(args.port)
        sys.exit(0)
    
    else:
        if not args.addr:
            print('server ip not defined, invalid parameters.', file=sys.stderr)
            sys.exit(1)
        if not args.file:
            log.error('filename not defined, invalid parameters.', file=sys.stderr)
            sys.exit(1)

        client_mode(args.addr, args.port, args.file)
        sys.exit(0)


if __name__ == '__main__':
    main()
    sys.exit(0)