4
\$\begingroup\$

To learn network programming in Python, I have written a small piece of code which is constituted of 2 modules, a client and a server. The server listens to a specific port provided on command line, then the client connects to the server, and send a file to it.

The client sends the size of file at first, then the content of file. During data transmission, both side of socket update a SHA256 object, and when transmission is completed, they calculate the hash digests and print them on screen.

I know my code is still far from being a robust implementation, so I am looking for your opinions. Any criticism is welcome. Thank you.

Module bigfile_server.py:

# -*- encoding: utf-8 -*-
import sys
import struct
import getopt
import socket
import hashlib
from datetime import datetime

FILE_BUFFER_SIZE = 524288

def usage():
    print('Usage: bigfile_server.py <SERVER_PORT>')
    print('SERVER_PORT: Port to which server will listen.')

def random_filename():
    dt_now = datetime.now()
    return dt_now.strftime('%Y%m%d%H%M%S%f')

def readn(sock, count):
    data = b''
    while len(data) < count:
        packet = sock.recv(count - len(data))
        if packet == '':
            return ''
        data += packet
    return data

if __name__ == '__main__':

    server_port = ''
    opts, args = getopt.gnu_getopt(sys.argv[1:], 'p:', ['port='])
    for opt, arg in opts:
        if opt in ('-p', '--port'):
            server_port = arg

    if server_port == '':
        print('Server port missing.', file=sys.stderr)
        usage()
        sys.exit(1)

    if not server_port.isdecimal():
        print('Server port contains invalid characters.', file=sys.stderr)
        sys.exit(2)

    print('Launching bigfile server.')
    serv_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        serv_sock.bind(('', int(server_port)))
        serv_sock.listen(5)
    except socket.error as e:
        print('Failed to launch server:', e)
        sys.exit(3)
    else:
        print('Server launched, waiting for new connection.')

    try:
        clnt_sock, clnt_addr = serv_sock.accept()
    except socket.error as e:
        print ('Failed to accept new connection:', e)
        sys.exit(3)
    else:
        print('New connection from:', clnt_sock)

    size_buff = readn(clnt_sock, 4)
    if size_buff == '':
        print('Failed to receive file size.', file=sys.stderr)
        clnt_sock.close()
        serv_sock.close()
        sys.exit(3)

    size_unpacked = struct.unpack('!I', size_buff)
    file_size = size_unpacked[0]
    print('Will receive file of size', file_size, 'bytes.')

    hash_algo = hashlib.sha256()

    filename = random_filename()
    try:
        with open(filename, 'wb') as file_handle:
            while file_size > 0:
                buffer = clnt_sock.recv(FILE_BUFFER_SIZE)
                print(len(buffer), 'bytes received.')
                if buffer == '':
                    print('End of transmission.')
                    break
                hash_algo.update(buffer)
                file_handle.write(buffer)
                file_size -= len(buffer)
            if file_size > 0:
                print('Failed to receive file,', file_size, 'more bytes to go.')
    except socket.error as e:
        print('Failed to receive data:', e, file=sys.stderr)
        clnt_sock.close()
        serv_sock.close()
        sys.exit(3)
    except IOError as e:
        print('Failed to write file:', e, file=sys.stderr)
        clnt_sock.close()
        serv_sock.close()
        sys.exit(3)
    else:
        print('File transmission completed.')

    clnt_sock.shutdown(socket.SHUT_RD)
    clnt_sock.close()
    serv_sock.close()
    print('Server shutdown.')
    print('SHA256 digest:', hash_algo.hexdigest())

    sys.exit(0)

Module bigfile_client.py:

# -*- encoding: utf-8 -*-

import os
import sys
import struct
import getopt
import socket
import hashlib

FILE_BUFFER_SIZE = 524288

def usage():
    print('Usage: bigfile_client.py <ARGUMENTS>')
    print('ARGUMENTS:')
    print('[-f|--file]: Source file to send.')
    print('[-h|--host]: Server address.')
    print('[-p|--port]: Server port number.')

if __name__ == '__main__':

    source_file, server_addr, server_port = '', '', ''
    opts, args = getopt.gnu_getopt(sys.argv[1:], 'f:h:p:', ['file=', 'host=', 'port='])
    for opt, arg in opts:
        if opt in ('-f', '--file'):
            source_file = arg
        elif opt in ('-h', '--host'):
            server_addr = arg
        elif opt in ('-p', '--port'):
            server_port = arg

    if source_file == '':
        print('Source file missing.', file=sys.stderr)
        usage()
        sys.exit(1)
    elif server_addr == '':
        print('Server address missing.', file=sys.stderr)
        usage()
        sys.exit(1)
    elif server_port == '':
        print('Server port number missing.', file=sys.stderr)
        usage()
        sys.exit(1)

    if not os.path.isfile(source_file):
        print('Source file cannot be found.', file=sys.stderr)
        sys.exit(2)
    elif not server_port.isdecimal():
        print('Server port number contains invalid characters.', file=sys.stderr)
        sys.exit(2)

    file_size = os.path.getsize(source_file)
    print('Sending file {0} to {1}:{2}.'.format(source_file, server_addr, server_port))
    print('Source file size:', file_size, 'bytes.')

    print('Connecting to remote server.')
    conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        conn.connect((server_addr, int(server_port)))
    except socket.error as e:
        print('Failed to connect to server:', e)
        sys.exit(3)
    else:
        print('Connection established.')

    print('Sending file size to remote server.')
    buffer = b''
    buffer = struct.pack('!I', file_size)
    print('File size packed into binary format:', buffer)

    try:
        conn.sendall(buffer)
    except socket.error as e:
        print('Failed to send file size:', e)
        sys.exit(3)
    else:
        print('File size sent.')

    hash_algo = hashlib.sha256()

    print('Start to send file content.')
    try:
        with open(source_file, 'rb') as file_handle:
            buffer = file_handle.read(FILE_BUFFER_SIZE)
            while len(buffer) > 0:
                conn.sendall(buffer)
                hash_algo.update(buffer)
                buffer = file_handle.read(FILE_BUFFER_SIZE)        
    except IOError as e:
        print('Failed to open source file', source_file, ':', e, file=sys.stderr)
        sys.exit(3)

    conn.shutdown(socket.SHUT_WR)
    conn.close()
    print('File sent, connection closed.')
    print('SHA256 digest:', hash_algo.hexdigest())

    sys.exit(0)
\$\endgroup\$

1 Answer 1

1
\$\begingroup\$

If I’m reading this correctly - I did not test it, the file cannot be larger than a 4 byte integer of bytes which is only a 4gb file, which in 2018 I wouldn’t have considered large since it fits in memory for most laptops.

I love that you are using getopts!

What I’m not loving is that my port can be a decimal (3.14), negative number, or even larger than the 70,000. Port 0 also means something special to the OS (and internet standards)

You can also shorten a lot of the args logic. For example an empty string is evaluated to false.

If not string:
    What to do if the string was empty.

I think the server sending the hash value to the client to confirm that the file had been uploaded/downloaded/transferred correctly would also make for a better user experience.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.