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)
