import socket
import time
import sha
import struct
import StringIO
import binascii

import pygtk
pygtk.require('2.0')
import gtk, gobject


_MTU = 481
_HEADER_LEN = 31
_MAGIC = 0xbaea4304
_TTL = 120 # 2 minutes

# Message segment packet types
_SEGMENT_TYPE_DATA = 0
_SEGMENT_TYPE_RETRANSMIT = 1

def _stringify_sha(sha):
	print_sha = ""
	for char in sha:
		print_sha = print_sha + binascii.b2a_hex(char)
	return print_sha

def _sha_data(data):
	sha_hash = sha.new()
	sha_hash.update(data)
	return sha_hash.digest()

class MessageSegment(object):
	# 4: magic (0xbaea4304)
	# 1: type
	# 2: segment number
	# 2: total segments
	# 2: message sequence number
	#20: total data sha1
	_HEADER_TEMPLATE = "! IbHHH20s"

	def _new_from_parts(self, msg_seq_num, segno, total_segs, data, master_sha):
		"""Construct a new message segment from individual attributes."""
		if not data:
			raise ValueError("Must have valid data.")
		if segno > 65535:
			raise ValueError("Segment number cannot be more than 65535.")
		if segno < 1:
			raise ValueError("Segment number must be greater than zero.")
		if total_segs > 65535:
			raise ValueError("Message cannot have more than 65535 segments.")
		if total_segs < 1:
			raise ValueError("Message must have at least one segment.")
		if msg_seq_num < 1:
			raise ValueError("Message sequence number must be greater than 0.")
		self._stime = time.time()
		self._data = data
		self._data_len = len(data)
		self._master_sha = master_sha
		self._segno = segno
		self._total_segs = total_segs
		self._msg_seq_num = msg_seq_num
		self._addr = None
		self._type = _SEGMENT_TYPE_DATA

		# Make the header
		self._header = struct.pack(self._HEADER_TEMPLATE, _MAGIC, self._type,
				self._segno, self._total_segs, self._msg_seq_num, self._master_sha)

	def _new_from_data(self, addr, data):
		"""Verify and construct a new message segment from network data."""
		if len(data) < _HEADER_LEN + 1:
			raise ValueError("Segment is less then minimum required length")
		stream = StringIO.StringIO(data)
		self._stime = None
		self._addr = addr

		# Determine and verify the length of included data
		stream.seek(0, 2)
		header_size = struct.calcsize(self._HEADER_TEMPLATE)
		self._data_len = stream.tell() - header_size
		stream.seek(0)

		# Read the header attributes
		(magic, seg_type, segno, total_segs, msg_seq_num, master_sha) = struct.unpack(self._HEADER_TEMPLATE,
				stream.read(header_size))

		# Sanity checks on the message attributes
		if seg_type != _SEGMENT_TYPE_DATA and seg_type != _SEGMENT_TYPE_RETRANSMIT:
			raise ValueError("Segment has invalid type.")
		if seg_type == _SEGMENT_TYPE_RETRANSMIT:
			if segno != 1 or total_segs != 1:
				raise ValueError("Retransmission request messages must have only one segment.")
		if magic != _MAGIC:
			raise ValueError("Segment does not have the correct magic.")
		if self._data_len < 1:
			raise ValueError("Segment must have some data.")
		if self._data_len > _MTU:
			raise ValueError("Data length must not be larger than the MTU (%s)." % _MTU)
		if segno < 1:
			raise ValueError("Segment number must be greater than 0.")
		if segno > total_segs:
			raise ValueError("Segment number cannot be larger than message segment total.")
		if total_segs < 1:
			raise ValueError("Message must have at least one segment.")
		if msg_seq_num < 1:
			raise ValueError("Message sequence number must be greater than 0.")

		self._type = seg_type
		self._segno = segno
		self._total_segs = total_segs
		self._msg_seq_num = msg_seq_num
		self._master_sha = master_sha

		# Reconstruct the data
		self._data = struct.unpack("! %ds" % self._data_len, stream.read(self._data_len))[0]

	def new_from_parts(msg_seq_num, segno, total_segs, data, master_sha):
		"""Static constructor for creation from individual attributes."""
		segment = MessageSegment()
		segment._new_from_parts(msg_seq_num, segno, total_segs, data, master_sha)
		return segment
	new_from_parts = staticmethod(new_from_parts)

	def new_from_data(addr, data):
		"""Static constructor for creation from a packed data stream."""
		segment = MessageSegment()
		segment._new_from_data(addr, data)
		return segment
	new_from_data = staticmethod(new_from_data)

	def stime(self):
		return self._stime

	def addr(self):
		return self._addr

	def segment_number(self):
		return self._segno

	def total_segments(self):
		return self._total_segs

	def message_sequence_number(self):
		return self._msg_seq_num

	def data(self):
		return self._data

	def master_sha(self):
		return self._master_sha

	def segment_type(self):
		return self._type

	def segment(self):
		"""Return a correctly formatted message that can be immediately sent."""
		return self._header + self._data

class MostlyReliablePipe(object):
	"""Implement Mostly-Reliable UDP.  We don't actually care about guaranteeing
	delivery or receipt, just a better effort than no effort at all."""

	def __init__(self, local_addr, remote_addr, port, data_cb, user_data=None):
		self._local_addr = local_addr
		self._remote_addr = remote_addr
		self._port = port
		self._data_cb = data_cb
		self._user_data = user_data
		self._started = False
		self._worker = 0
		self._seq_counter = 0

		self._outgoing = []
		self._sent = []

		self._incoming = {}  # (message sha, # of segments) -> [segment1, segment2, ...]

		self._setup_listener()
		self._setup_sender()

	def _setup_sender(self):
		"""Setup the send socket for multicast."""
		self._send_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
		# Make the socket multicast-aware, and set TTL.
		self._send_sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 20) # Change TTL (=20) to suit

	def _setup_listener(self):
		"""Set up the listener socket for multicast traffic."""
		# Listener socket
		self._listen_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

		# Set some options to make it multicast-friendly
		self._listen_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
		self._listen_sock.setsockopt(socket.SOL_IP, socket.IP_MULTICAST_TTL, 20)
		self._listen_sock.setsockopt(socket.SOL_IP, socket.IP_MULTICAST_LOOP, 1)

	def start(self):
		"""Let the listener socket start listening for network data."""
		# Set some more multicast options
		self._listen_sock.bind((self._local_addr, self._port))  # Bind to all interfaces
		self._listen_sock.settimeout(2)
		intf = socket.gethostbyname(socket.gethostname())
		self._listen_sock.setsockopt(socket.SOL_IP, socket.IP_MULTICAST_IF,
				socket.inet_aton(intf) + socket.inet_aton('0.0.0.0'))
		self._listen_sock.setsockopt(socket.SOL_IP, socket.IP_ADD_MEMBERSHIP,
				socket.inet_aton(self._remote_addr) + socket.inet_aton('0.0.0.0'))

		# Watch the listener socket for data
		gobject.io_add_watch(self._listen_sock, gobject.IO_IN, self._handle_incoming_data)
		gobject.timeout_add(120000, self._segment_ttl_worker)

		self._started = True

	def _segment_ttl_worker(self):
		"""Cull already-sent message segments that are past their TTL."""
		now = time.time()
		for segment in self._sent[:]:
			if segment.stime() < now - _MSG_TTL:
				self._sent.remove(segment)
		return True

	def _dispatch_message(self, addr, message):
		"""Send complete message data to the owner's data callback."""
		self._data_cb(addr, message, self._user_data)

	def _process_incoming_data(self, segment):
		"""Handle a new message segment.  First checks if there is only one
		segment to the message, and if the checksum from the header matches
		that computed from the data, dispatches it.  Otherwise, it adds the
		new segment to the list of other segments for that message, and 
		checks to see if the message is complete.  If all segments are present,
		the message is reassembled and dispatched."""

		string_sha = _stringify_sha(segment.master_sha())
		nsegs = segment.total_segments()
		addr = segment.addr()
		segno = segment.segment_number()

		# Short-circuit single-segment messages
		if segno == 1 and nsegs == 1:
			# Ensure the header's master sha actually equals the data's sha
			if string_sha == _stringify_sha(_sha_data(segment.data())):
				self._dispatch_message(addr, segment.data())
				return

		# Otherwise, track the new segment
		msg_seq_num = segment.message_sequence_number()
		msg_key = (addr[0], msg_seq_num, string_sha, nsegs)
		if not self._incoming.has_key(msg_key):
			self._incoming[msg_key] = {}

		# Look for a dupe, and if so, drop the new segment
		if self._incoming[msg_key].has_key(segno):
			return
		self._incoming[msg_key][segno] = segment

		# Dispatch the message if all segments are present and the sha is correct
		if len(self._incoming[msg_key]) == nsegs:
			all_data = ''
			for i in range(1, nsegs + 1):
				all_data = all_data + self._incoming[msg_key][i].data()
			if string_sha == _stringify_sha(_sha_data(all_data)):
				self._dispatch_message(addr, all_data)
			del self._incoming[msg_key]

	def _process_retransmit_request(self, segment):
		"""Validate and process a retransmission request."""
		# Retransmission data format:
		#  2: message sequence number
		# 20: total data sha1
		#  2: segment number
		data = segment.data()
		if len(data) != 22:
			print "Bad retransmission request message format."
		# Native byte-order since the receive bits already unpacked it for us
		template = "@ H20sH"
		(msg_seq_num, master_sha, segno) = struct.unpack(template, data)

	def _handle_incoming_data(self, source, condition):
		"""Handle incoming network data by making a message segment out of it
		sending it off to the processing function."""
		if not (condition & gobject.IO_IN):
			return True
		msg = {}
		data, addr = source.recvfrom(_MTU + _HEADER_LEN)
		try:
			segment = MessageSegment.new_from_data(addr, data)
			if segment.segment_type() == _SEGMENT_TYPE_DATA:
				self._process_incoming_data(segment)
			elif segment.segment_type() == _SEGMENT_TYPE_RETRANSMIT:
				self._process_retransmit_request(segment)
		except ValueError, exc:
			print "Bad segment: %s" % exc
		return True

	def send(self, data):
		"""Break data up into chunks and queue for later transmission."""
		if not self._started:
			raise Exception("Can't send anything until started!")

		self._seq_counter = self._seq_counter + 1
		if self._seq_counter > 65535:
			self._seq_counter = 1

		# Pack the data into network byte order
		template = "! %ds" % len(data)
		data = struct.pack(template, data)
		master_sha = _sha_data(data)

		# Split up the data into segments
		left = length = len(data)
		nmessages = length / _MTU
		if length % _MTU > 0:
			nmessages = nmessages + 1
		msg_num = 1
		while left > 0:
			msg = MessageSegment.new_from_parts(self._seq_counter, msg_num,
					nmessages, data[:_MTU], master_sha)
			self._outgoing.append(msg)
			msg_num = msg_num + 1
			data = data[_MTU:]
			left = left - _MTU
		if len(self._outgoing) > 0 and self._worker == 0:
			self._worker = gobject.idle_add(self._send_worker)

	def _send_worker(self):
		"""Send all queued segments that have yet to be transmitted."""
		self._worker = 0
		for segment in self._outgoing:
			data = segment.segment()
			self._send_sock.sendto(data, (self._remote_addr, self._port))
		self._sent = self._outgoing
		self._outgoing = []
		return False


def got_data(addr, data, user_data=None):
	print "Data (%s): %s" % (addr, data)

def main():
	pipe = MostlyReliablePipe('', '224.0.0.222', 2293, got_data)
	pipe.start()
	pipe.send('The quick brown fox jumps over the lazy dog')
	gtk.main()


if __name__ == "__main__":
	main()