Building my own TCP client and sender for networking class!

receiver.py

#!/usr/bin/env python 

""" 
Receives a file via mock TCP over UDP
""" 
import select
import socket 
import sys 
import struct
import hashlib
import binascii
import time
import datetime

class Receiver:
	def __init__(self, listening_port = 3000, remote_IP=None,remote_port=8000, host='localhost', port=50000, logfile="receiver.log", filename="new_file"):
		self.listening_port = int(listening_port)
		self.remote_port = int(remote_port)
		self.host = host
		self.port = port
		self.size = 1024 
		self.remote_IP = remote_IP
		self.backlog=5
		if logfile != "stdout":
			self.logfile=open(logfile,"w")
		else:
			self.logfile = "stdout"
		self.new_file=open(filename,"wb")

		self.ack_number = 0
		# was the last packet acknowledged
		self.ACKed = False
		self.run()

	def run(self):
		# a listening socket for the file to be transferred
		file_receiver = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		file_receiver.bind((self.host,self.listening_port))
		file_receiver.listen(self.backlog)

		# open socket to send ACKs
		ack_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 
		print self.remote_port
		ack_socket.connect((self.remote_IP,self.remote_port)) 

		# start initial conditions for select
		output = [ack_socket]
		input = [file_receiver]
		running = 1

		# start select loop
		while running: 
			inputready, outputready, exceptread = select.select(input, output, [])
			for s in inputready:
				if s == file_receiver:
					# accept file sender handshake
					file_listener, address = file_receiver.accept()
					# add to inputs
					input.append(file_listener)
			
				# file listening socket
				if s == file_listener:
					# receive TCP
					data = file_listener.recv(20)
					if data:
						struct_format = '=iisis'
						# unpack data
						try:
							data_unpacked = struct.unpack(struct_format,data)
						except:
							break
						remote_port = data_unpacked[0]
						ack_number = data_unpacked[1]
						# if duplicate packet
						if ack_number == self.ack_number:
							# stop process
							break
						internet_checksum = data_unpacked[2]
						fin = data_unpacked[3] 
						#msg = binascii.b2a_qp(data_unpacked[4])
						msg = data_unpacked[4]

						#print "remote port %s ack_number %s internet_checksum %s fin %s msg %s" %(str(remote_port), str(ack_number),str(internet_checksum),str(fin),str(msg))
						if internet_checksum == "error":
							raise Exception("This file does not exist.")

						# tranmission is over
						if fin == 1:
							# log the end of the transmission
							self.write_log(fin=1, ack_number=ack_number)
							# tell console the transmission is completed
							print "Transmission completed"
							# remove the file listener socket
							input.remove(file_listener)
							# remove the ack socket
							output.remove(ack_socket)
							# close the file listener
							file_listener.close()
							ack_socket.close()
							self.new_file.close()
							sys.exit()

						our_checksum = self.gen_checksum(remote_port=remote_port, ack_number=ack_number,fin=fin, data=msg)[0]
		

						# if checksum is the same
						if our_checksum == internet_checksum:
							self.ack_number = self.ack_number + 1
							self.ACKed = True
							# write data to file
							self.new_file.write(msg)
							self.write_log(fin=0, ack_number=self.ack_number)
						# if checksum is incorrect
						else:
							pass
					# no data
					else:
						pass
				break

			for s in outputready:
				if self.ACKed:
					s.send(str(self.ack_number))
					self.ACKed = False
				break

		# close file
		self.new_file.close()

	def write_log(self, fin=0, ack_number=0):
		"""
		Write to logfile

		args:
			fin: The FIN flag. default to 0 since the transimission ends once (if done successfully).
		"""
		self.time_packet_sent = time.time()
		timestamp = datetime.datetime.fromtimestamp(self.time_packet_sent).strftime('%Y-%m-%d %H:%M:%S')
		destination = self.host + ":" + str(self.listening_port)
		source = self.remote_port
		log_string = "Timestamp %s, Source %s, Destination %s, Ack %s \n" %(timestamp,str(self.remote_port), str(self.listening_port), str(self.ack_number))
		if self.logfile != "stdout":
			self.logfile.write(log_string)
		else:
			print log_string

	def gen_checksum(self, **kwargs):
		"""
		generate a checksum with a MD5 hash
		"""

		checksum_data_list = [str(kwargs['remote_port']), str(kwargs['ack_number']), str(kwargs['fin']), str(kwargs['data'])]
		checksum_data = " ".join(checksum_data_list)
		return hashlib.md5(checksum_data).hexdigest()

print sys.argv
if len(sys.argv) < 6:
	raise Exception("Too few commandline arguments.")

filename = sys.argv[1] 
listening_port = sys.argv[2]
remote_IP = sys.argv[3]
remote_port = sys.argv[4]
log_filename = sys.argv[5]

Receiver(filename=filename, listening_port=listening_port, remote_IP=remote_IP, remote_port=remote_port, logfile=log_filename)

sender.py

#!/usr/bin/env python 

""" 
Sends a file over mock TCP over UDP
""" 

import select 
import socket 
import sys 
import struct
import hashlib
import datetime
import time

class Sender:
	def __init__(self,listening_port=8000,remote_IP=None,remote_port = 3000, filename="test-file.txt", window_size=1, logfile="sender.log", ack_port_number=None):

		#TODO: replace with args
		self.host = '' 
		self.listening_port = int(listening_port)
		self.remote_port = int(remote_port)
		self.backlog = 5 
		self.infile = filename
		self.window_size = window_size
		self.remote_IP = remote_IP
		if logfile !="stdout":
			self.logfile=open(logfile, "w")
		else:
			self.logfile = "stdout"
		# last packet acked
		self.ack_number = 0
		# last packet number attempted to send
		self.last_ack = -1
		self.sample_RTT = 0
		self.dev_RTT = 0
		self.estimated_RTT = 0
		self.sent_segments = 0
		self.seq_number = 0
		self.fin_flag = 0
		self.internet_checksum="x"
		self.segments_retransmitted = 0
		if ack_port_number:
			self.ack_port_number = int(ack_port_number)
		try:
			self.file = open(self.infile, "rb")
			self.is_file = True
		except:
			self.is_file = False
			raise Exception("This file does not exist.")
		self.run()

	def run(self):

		# bind a socket to listen for ACKs
		ack_receiver = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 
		ack_receiver.bind((self.remote_IP,self.ack_port_number)) 
		ack_receiver.listen(self.backlog) 

		file_sender = None
		output = []
		error = []
		input = [ack_receiver,sys.stdin] 
		packet_number = 0
		running = 1 

		while running: 
			inputready,outputready,exceptready = select.select(input,output,error) 
			for s in inputready:
				if s == ack_receiver: 
					# handle the ack socket 
					receiver, address = ack_receiver.accept() 
					input.append(receiver) 
					# connect a socket to send the file
					file_sender = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
					file_sender.connect((self.host, self.remote_port))
					output.append(file_sender)

				elif s == sys.stdin: 
					# handle standard input 
					junk = sys.stdin.readline() 
					running = 0 

				else:
					ack = s.recv(self.window_size)
					if ack:
						# update packet number to send
						self.ack_number = self.ack_number + 1
						self.seq_number = (self.window_size * self.ack_number)
						self.sample_RTT = time.time() - self.time_packet_sent
					else:
						pass

			for s in outputready:
				# if the socket is the file sending socket
				if s == file_sender:
					# file does exist
					if self.is_file:
						header_length = 20
						receive_window = self.window_size

						if self.last_ack < self.ack_number:

							self.data = self.file.read(self.window_size)

							if self.data:
								# store last ack
								self.last_ack = self.ack_number
								# calculate checksum
								self.internet_checksum = str(self.gen_checksum())[0]
								# pack header
								header = self.pack_header()				
								# send header
								file_sender.send(header)
								# write to log
								self.write_log()

							# no data, send fin flag and end transmission
							else:
								self.write_log(fin=1)
								self.data = " "
								self.internet_checksum = "x"
								header = self.pack_header(fin=1)	
								# send last header
								file_sender.send(header)
								print "Delivery completed successfully"
								print "Total bytes sent %i" %(self.seq_number + self.segments_retransmitted)
								print "Segments sent %i" %(self.ack_number)
								print "Segments retransmitted %i" %(self.segments_retransmitted)


								# close socket
								file_sender.close()
								receiver.close()
								ack_receiver.close()
								if self.logfile != "stdout":
									self.logfile.close()
								sys.exit()

							# if no ack
						else:
							# if packet has timed out
							if time.time() > self.calculate_estimated_RTT():
								# attempt to send packet again
								self.internet_checksum = str(self.gen_checksum())[0]
								# pack header
								header = self.pack_header()				
								# send header
								file_sender.send(header)
								# write to log
								self.write_log()
								# add to segments retransmitted


					# file does not exist		
					else:
						raise Exception("No such file.")
						self.internet_checksum = "error"
						# pack header
						header = self.pack_header(fin=1)			
						# send header
						file_sender.send(header)


	def pack_header(self, fin=0):
		"""
		Pack header struct. Fin defaults to 0 since it is only 1 on the last packet.

		args:
			fin: fin flag to signal if the transmission is the last transmission or not
		"""
		struct_format = '=iisis'
		header = struct.pack(struct_format,self.remote_port,self.ack_number,self.internet_checksum, fin, self.data)	
		return header			

	def write_log(self, fin=0):
		"""
		Write to logfile

		args:
			fin: The FIN flag. default to 0 since the transimission ends once (if done successfully).
		"""
		self.time_packet_sent = time.time()
		timestamp = datetime.datetime.fromtimestamp(self.time_packet_sent).strftime('%Y-%m-%d %H:%M:%S')
		destination = self.host + ":" + str(self.listening_port)
		source = self.infile
		ACK_flag = 1 
		SYN_flag = self.seq_number
		FIN_flag = 0
		self.estimated_RTT = self.calculate_estimated_RTT()
		log_string = "Timestamp: %s Source: %s, Destination: %s, Sequence Number: %i, Estimated RTT: %i, ACK number: %i, Fin flag:%i \n" %(self.time_packet_sent, source, destination, self.seq_number, self.estimated_RTT, self.ack_number, FIN_flag) 
		if self.logfile != "stdout":
			self.logfile.write(log_string)
		else:
			print log_string

	def gen_checksum(self):
		"""
		generate a checksum with a MD5 hash
		"""
		checksum_data_list = [str(self.remote_port), str(self.ack_number), str(self.fin_flag), str(self.data)]
		checksum_data = " ".join(checksum_data_list)
		return hashlib.md5(checksum_data).hexdigest()


	def calculate_estimated_RTT(self):
		"""
		Calculate an estimated roundtrip time
		"""
		alpha = 0.125
		estimate = self.estimated_RTT*(1-alpha) + self.sample_RTT*(alpha)
		return estimate

	def calculate_dev_RTT(self):
		"""
		calcualte RTT deviation
		"""
		beta = 0.25
		( 1- beta) * self.dev_RTT + (beta * abs(self.sampleRTT - self.estimated_RTT))

	def timeout(self):
		"""
		Set timeout
		"""
		self.timeout = self.estimated_RTT + ( 4  * self.dev_RTT)

print sys.argv
if len(sys.argv) < 7:
	raise Exception("Too few commandline arguments.")

filename = sys.argv[1]
remote_IP = sys.argv[2]
remote_port = sys.argv[3]
ack_port_number = sys.argv[4]
window_size= sys.argv[5]
log_filename = sys.argv[6]

Sender(filename=filename, remote_IP=remote_IP, remote_port=remote_port,
	ack_port_number=ack_port_number,window_size=window_size, logfile=log_filename)