#!/usr/local/bin/python2.4
# 
# Copyright 2005 John-Mark Gurney
# All rights reserved.
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
# 

import copy
import divert
import dpkt
import getopt
import math
import time
import sys

def getkey(ip):
	tcp = ip.data
	return (ip.src, tcp.sport, ip.dst, tcp.dport)

def calcackbytes(acka, ackb):
	r = ackb - acka
	if r < -2**32 + 2**31:	# 1 gig window size for SACK
		r = ackb - acka + 2 ** 32
	return r

def cmpiptcp(a, b):
	# compare ip
	r = cmp(getkey(a), getkey(b))
	if r != 0:
		return r

	# compare tcp
	f = lambda x: (x.data.seq, x.data.ack, x.data.data)
	return cmp(f(a), f(b))

class acklimit(divert.Divert):
	def __init__(self, port, bytespersec, **kwargs):
		self.bytespersec = bytespersec
		self.cons = {}
		self.pending = []
		self.updatelt = True
		self.lasttime = 0
		self.waittime = 24*60*60
		apply(divert.Divert.__init__, (self, port), kwargs)

	def cansend(self, bytes):
		curt = time.time()
		if curt - self.lasttime >= 0:
			return True
		else:
			self.waittime = self.lasttime - curt
			return False

	def newtime(self, bytes):
		if self.updatelt:
			bt = time.time()
		else:
			bt = self.lasttime
		self.lasttime = bt + float(bytes) / self.bytespersec

	def dotcpack(self, ip, addr):
		assert isinstance(ip.data, dpkt.tcp.TCP) and ip.data.flags & dpkt.tcp.TH_ACK

		tcp = ip.data
		key = getkey(ip)

		try:
			lastack, pqueue = self.cons[key]
			ackbytes = calcackbytes(lastack, tcp.ack)
			if ackbytes <= 0:
				self.sendpkt(str(ip), addr)
			else:
				mat = False
				for b, addr in pqueue:
					if cmpiptcp(b, ip) == 0:
						mat = True
						break
				if not mat:
					i = (ip, addr)
					self.pending.append(i)
					pqueue.append(i)
		except KeyError:
			# no connection, create one
			self.cons[key] = (ip.data.ack, [])
			self.sendpkt(str(ip), addr)

	def checkpending(self):
		while self.pending:
			ip, addr = self.pending[0]
			tcp = ip.data
			key = getkey(ip)
			lastack, pqueue = self.cons[key]
			ackbytes = calcackbytes(lastack, tcp.ack)
			if ackbytes <= 0:
				self.sendpkt(str(ip), addr)
			else:
				if not self.cansend(ackbytes):
					break
				self.newtime(ackbytes)
				self.sendpkt(str(ip))
				self.updatelt = False
			self.cons[key] = (ip.data.ack, pqueue)
			del self.pending[0]
			del pqueue[0]
		if not self.pending:
			self.waittime = 24*60*60
			self.updatelt = True

	def sendpkt(self, buf, addr = None):
		ret = divert.Divert.sendpkt(self, buf, addr)
		if 0:
			print 'sp, len:', len(buf), 'send:', ret

		return ret

	def droppending(self, pqueue):
		for i in pqueue:
			self.pending.remove(i)

	def handlepkt(self, buf, addr):
		ip = dpkt.ip.IP(buf)
		echopkt = True
		if isinstance(ip.data, dpkt.tcp.TCP):
			if ip.data.flags & dpkt.tcp.TH_FIN:
				try:
					key = getkey(ip)
					lastack, pqueue = self.cons[key]
					self.droppending(pqueue)
					del self.cons[key]
				except KeyError:
					pass
			elif ip.data.flags & dpkt.tcp.TH_ACK:
				self.dotcpack(ip, addr)
				echopkt = False

		if echopkt:
			self.sendpkt(buf, addr)

def usage():
	sys.stderr.write('''Usage: acklimit -p <divertport> -b <bytespersec>\n''')

def main():
	try:
		opts, args = getopt.getopt(sys.argv[1:], "b:hp:", ["help", ])
	except getopt.GetoptError:
		# print help information and exit:
		usage()
		sys.exit(2)

	dvrtprt = None
	bytespersec = None
	for o, a in opts:
		if o in ("-h", "--help"):
			usage()
			sys.exit()
		elif o in ("-p", ):
			dvrtprt = int(a)
		elif o in ("-b", ):
			bytespersec = int(a)

	if not dvrtprt or not bytespersec:
		usage()
		sys.exit(2)

	d = acklimit(dvrtprt, bytespersec, block = False)
	while True:
		d.checkpending()
		d.checkandhandle(d.waittime)

if __name__ == '__main__':
	main()