simple UDP packet loss tester

ploss.py

"""

Usage:
(1) on host A:
	python ploss.py recv <port>
(2) on host B:
	python ploss.py send <A's hostname> <port from step 1>
(3) watch output on host B

"""
import socket
import sys
import time
import random
import struct
import os
import shutil

MAX_RECV_SIZE = 1024*16

def safe_div(a,b):
	if b == 0:
		return 0
	else:
		return a/b
		
class PacketInfo(object):
	def __init__(self, id, size, send_time, recv_time):
		self.id = id
		self.size = size
		self.send_time = send_time
		self.recv_time = recv_time

class ClientState(object):
	def __init__(self, addr, s):
		self.addr = addr
		self.highest_id = -1
		self.received = []
		self.socket = s
		self.last_update = 0
		self.total_size_received = 0
		self.start_time = time.time()
		self.ooo = 0
		
	def on_recv(self, packet_time, data):
		type = data[8]
		# Data packet
		if type == 'd':
			#print '%s Received a data packet' % self
			un1,un2,id,send_time = struct.unpack("!QcQd", data[:(8+1+8+8)])
			if id != self.highest_id+1:
				self.ooo += 1
			self.received.append(PacketInfo(id, len(data), send_time,packet_time))
			self.total_size_received += len(data)
			self.highest_id = max(id, self.highest_id)
		# Report request
		elif type == 'r':
			#print '%s Got report request of %d bytes' % (self,len(data))
			un1,un2,timestamp = struct.unpack("!Qcd", data[:(8+1+8)])
			dt = time.time()-self.start_time
			report = struct.pack("!QQQdddd",
									self.highest_id+1,
									len(self.received),
									self.ooo,
									loss_percent(self.highest_id+1, len(self.received)),
									safe_div(len(self.received),dt),
									safe_div(self.total_size_received,dt),
									timestamp)
			#print 'Sending %d byte report request to %s' % (len(report), repr(self.addr))
			self.socket.sendto(report, self.addr)
		else:
			print 'Unknown type: %s' % repr(type)
			

			
def loss_percent(known_sent, received):
	if known_sent == 0:
		return 0
	return 100.0 * float(known_sent - received)/known_sent

def format_byte_rate(rate):
	if rate >= 1024*1024:
		return '%.1f MB/s'%(rate/(1024*1024))
	elif rate >= 1024:
		return '%.1f kB/s'%(rate/(1024))
	else:
		return '%.1f B/s'%rate

if sys.argv[1] == 'recv':

	shutil.rmtree('client_data', True)
	time.sleep(0.2)
	os.mkdir('client_data')
		
	s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
	s.bind(('0.0.0.0', int(sys.argv[2])))
	clients = {}
		
	start = time.time()
		
	try:
		while True:
			data,fromaddr = s.recvfrom(MAX_RECV_SIZE)
			t = time.time()
			key = (fromaddr, data[:8])
			if key in clients:
				client = clients[key]
			else:
				client = ClientState(fromaddr, s)
				clients[key] = client
			client.on_recv(t, data)
	except KeyboardInterrupt:
		n=0
		for client in clients.values():
			with open('client_data/%d.csv'%n, 'wt') as f:
				for r in client.received:
					f.write('%d,%f,%f\n'%(r.size, r.send_time, r.recv_time))
			n+=1
		
elif sys.argv[1] == 'send':
	
	def rand_uint64():
		return random.randint(0, 2**64-1)
	
	server_addr = (sys.argv[2], int(sys.argv[3]))
	s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
	s.bind(('0.0.0.0', 0))
	stream = rand_uint64()
	id = 0
	interval = 0
	#rate = 22
	rate = 30
	while True:
		size = random.randint(0,1400)
		randomness = os.urandom(size)
		data = struct.pack("!QcQd", stream, 'd', id, time.time())+randomness
		s.sendto(data, server_addr)
		id += 1
		time.sleep(1.0/rate)
		interval += 1
		if interval >= rate:
			interval = 0
			request = struct.pack("!Qcd", stream, 'r', time.time())
			s.sendto(request, server_addr)
			s.settimeout(2.0)
			response = None
			try:
				response,from_addr = s.recvfrom(MAX_RECV_SIZE)
				response_time = time.time()
			except:
				print 'Timed out waiting for report'
				pass
			s.settimeout(None)
			if response is not None:
				report = struct.unpack("!QQQdddd",response)
				print "Sent=%d Recv'd=%d OOO=%d Loss=%.2f %.1f pkts/s %s Ping=%.1f ms" % (
					report[0],
					report[1],
					report[2],
					report[3],
					report[4],
					format_byte_rate(report[5]),
					1000.0 * (time.time()-report[6])
				)