Twisted TFTP Server with client map

tftp.py

# Anthony Clark
# 
# Simple TFTP Server in Twisted
# TODO timeouts/error handling
#
#
# Most TFTP server do not have a 
# way to change file root per client.
# This one does however using the client
# map below

import os, binascii

from twisted.internet.protocol import DatagramProtocol
from twisted.internet import reactor


toInt 	= lambda x: int(binascii.hexlify(x), 16)
toByte  = lambda x: binascii.unhexlify(hex(x)[2:].rjust(4, '0'))


# Simple TFTP Server
class TFTP(DatagramProtocol):
	
	# (host, port) to filename mapping
	current_clients = {}


	def __init__(self, host_map):
		self.host_map = host_map

	
	# File data generator
	def get_chunk(self, (host, port), blksize=512):
		while True:
			data = self.current_clients[(host,port)].read(blksize)
			if not data: break
			yield data


	# Send data
	def send_chunk(self, block_number, (host, port)):
		try:
			self.transport.write('\x00\x03%s%s' %
					(toByte(block_number), self.get_chunk((host,port)).next()), 
					(host, port))
		except:
			self.current_clients[(host, port)].close()
			del self.current_clients[(host,port)]
			pass


	# Receive packet
	def datagramReceived(self, data, (host, port)):
		# Get class mapping
		_class = self.host_map.get(host, '.')

		# READ Request
		if data.startswith('\x00\x01'):
			filename = _class + data.split('\x00')[1][1:]
			mode = data.split('\x00')[2]

			if not os.path.exists(filename):
				self.transport.write('\x00\x05\x00\x01no such file exists\x00', (host, port))
				return
			
			# file info
			fsize = os.path.getsize(filename)
			print host, "(", mode, ")", "wants", filename, "with size", fsize, "bytes"
			
			if fsize > 65535:
				print 'File too big:',fsize
				self.transport.write('\x00\x05\x00\x04illegal TFTP operation\x00', (host, port))
				return

			# open file
			self.current_clients[(host,port)] = open(filename, 'r')

			# Send first block of data
			self.send_chunk(1, (host, port))

		# ACK Request
		elif data.startswith('\x00\x04'):
			# read block from 2byte ack block
			new_block = toInt(data[2:][0]) << 8 | toInt(data[2:][1]) & 0xff
			self.send_chunk(new_block + 1, (host, port))



# Example client map
m = { '10.0.1.51' : 'class0/',
      '192.168.0.107': 'class0/',
      '127.0.0.1' : 'class0/'}

# listen
reactor.listenUDP(5000, TFTP(m))
reactor.run()