3
\$\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\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.