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)