import os
import nfqueue
import scapy
import signal

active_connexions = {}

def callback(packet):
	clientport1 = None
	ip = scapy.IP(packet.payload)
	if ip.proto == 6:
		tcp = ip.payload
		if tcp.flags == 17: # FIN, ACK -> fin de connexion
			if active_connexions.has_key(((ip.src, tcp.sport), (ip.dst, tcp.dport))):
				os.system('iptables -t nat -D PREROUTING -m udp -p udp -s %s --dport %s -j DNAT --to-destination %s' % (ip.src, active_connexions[((ip.src, tcp.sport), (ip.dst, tcp.dport))], ip.dst))
				del active_connexions[((ip.src, tcp.sport), (ip.dst, tcp.dport))]
				return (nfqueue.NF_ACCEPT, packet.nfmark, str(ip))

		p = str(tcp.payload)
		for l in p.split('\n'):
			if l[:10] == "Transport:":
				for args in l.strip().split(';'):
					if '=' in args:
						name, value = args.split('=')
						if name == 'client_port':
							clientport1, clientport2 = value.split('-')
							break
	if clientport1:
		active_connexions[((ip.src, tcp.sport), (ip.dst, tcp.dport))] = clientport1
		os.system('iptables -t nat -I PREROUTING -m udp -p udp -s %s --dport %s -j DNAT --to-destination %s' % (ip.src, clientport1, ip.dst))
		return (nfqueue.NF_ACCEPT, packet.nfmark, str(ip))
	return (nfqueue.NF_ACCEPT, packet.nfmark, str(ip))

def SIGINTHandler(signum, frame):
	for k, v in active_connexions.items():
		os.system('iptables -t nat -D PREROUTING -m udp -p udp -s %s --dport %s -j DNAT --to-destination %s' % (k[0][0], v,  k[1][0]))
	os.system('iptables -t filter -D FORWARD -m tcp -p tcp --sport 554 -j QUEUE')
	sys.exit(0)

if __name__ == '__main__':
	signal.signal(signal.SIGINT,SIGINTHandler)
	os.system('iptables -t filter -I FORWARD -m tcp -p tcp --sport 554 -j QUEUE')
	queue = nfqueue.Queue(False)
	qh = queue.createQueue(0, callback, nfqueue.NFQNL_COPY_PACKET)
	while True:
		queue.handlePacket()
	queue.destroyQueue(qh)