#!/usr/bin/env python2
import traceback, sys
from scapy.all import *
from tests.core import *
import hashlib
import hmac
import math
import socket

# Helper functions ###########################################################
    
def genhmac2(k1, k2, r1, r2):
    """Generate and return a HMAC-SHA1 with the concatenation of k1 and k2
    as key and the concatenation of r1 and r2 as message.

    k1, k2 are 64bits integers
    r1, r2 are 32bits integers
    Return a 160bits integer
    """
    key = xstr(k1).rjust(8,'\00') + xstr(k2).rjust(8,'\00')
    msg = xstr(r1).rjust(4,'\00') + xstr(r2).rjust(4,'\00') 
    return xlong(hmac.new(key, msg=msg, digestmod=hashlib.sha1).digest())
def genhmac(k1, k2, r1, r2):
    """Generate and return a HMAC-SHA1 with the concatenation of k1 and k2
    as key and the concatenation of r1 and r2 as message.

    k1, k2 are 64bits integers
    r1, r2 are 32bits integers
    Return a 160bits integer
    """
    key = xstr(k1).rjust(8,'\00') + xstr(k2).rjust(8,'\00')
    msg = xstr(r1).rjust(4,'\00') + xstr(r2).rjust(4,'\00') 
    return xlong(hmac.new(key, msg=msg, digestmod=hashlib.sha1).digest())

def key2tokenAndDSN(key):
    """Returns the token and dsn from a key
    Generate a simple SHA1 hash of the key

    key is a 64bits integer
    Token is a 32bits integer, dsn is a 64bits integer
    """
    import binascii
    keystr = struct.pack("!Q", key)
    h = hashlib.sha1(keystr.rjust(8,'\00'))
    shastr=h.digest() # binary
    #shastr = struct.pack("!IIIII", *struct.unpack("@IIIII",shastr)) #to net
    token, dsn = shastr[0:4], shastr[-8:]
    #print "raw: %s (len=%i)"%(shastr,len(shastr)) 
    #print "hex: %s"% binascii.hexlify(token), "%s"%binascii.hexlify(dsn)
    d1, d2 = struct.unpack("!II",dsn)
    token, dsn = (struct.unpack("!I",token)[0], (long(d2)<<32)+d1)
    #print "token: %x"% token
    #print "dsn: %x" % dsn
    return (token, dsn)

    
def getMpOption(tcp):
    """Return a generator of mptcp options from a scapy TCP() object"""
    for opt in tcp.options:
        if opt.kind == 30:
            yield opt.mptcp

def getMpSubkind(pkt, kind):
    """Return a generator of mptcp kind suboptions from pkt"""
    l4 = pkt.getlayer("TCP")
    for o in getMpOption(l4):
        if MPTCP_subtypes[o.subtype] == kind:
            yield (l4, o)

def checkAndGetMPOption(pkt, kind):
    """Return the first option of subkind kind in pkt
    If no such option exist, an exception is raised"""
    try:
        return getMpSubkind(pkt, kind).next()
    except StopIteration:
        raise Exception("MPTCP option of kind %s not found."%kind)
    
def genDSSChecksum(dsn, ssn, datalen, payload):
    """Generate a DSS Checksum on a data according to current state"""
    header = struct.pack("!QIHH", dsn, ssn, datalen, 0)
    return checksum(header+payload)

def getDataAckForPkt(s, sub, l4, plen, f=False):
    # should update the data_ack only if the map is relative to this
    # data segment.
    if sub["map"] and (l4.seq > sub["map"]["subseq"]+sub["startseq"]\
            and l4.seq < sub["map"]["subseq"]+sub["startseq"] + plen or f):
        if f and sub["map"]["subseq"] == 0: subseq = 1
        else: subseq = sub["map"]["subseq"]
        ret = sub["map"]["dsn"] \
                + (l4.seq -(subseq+sub["rem_startseq"])) \
                + plen
    else:
        ret = s["data_ack"]
    return ret

def get32bitSeq(seq64):
    seq32 = seq64 % (1 << 32) 
    return seq32

def kernelEstablishConn(t, mptcp, s, dst=None, dport=80):
    """Make the kernel establish a TCP connection (MPTCP if available in 
    the kernel) to dst"""
    self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    exchange = [mptcp.cap_syn, mptcp.cap_synack, mptcp.ack]
    self.sock.connect((dst, dport))

#############################################################################

class MPTCPTest(object):
    """MPTCP Test library

    Set of methods to test an MPTCP extension implementation. 
    It works together with the test framework.
    Each method uses and maintains a connection state. Every data useful for
    the packet generation should be stored in that object. 
    It also contains a model of the last received packet if there is one.
    Operations involving subflows are provided with a specific subflow state."""
    
    def __init__(self, tester=None, initstate=None):
        self.tester = tester
        tester.proto = self
        if initstate is None:
            initstate = MPTCPState()
        tester.state = initstate
        
        self.DefaultPacket = self.TCPPacket
        self.Ack = self.TCPPacket
        self.Push = self.TCPPacket

    def findProtoLayer(self, pkt):
        """Return an iterator on representations of proto components to
        consider"""
        mp=getMpOption(pkt.getlayer("TCP"))
        i=0
        for p in mp:
            i += 1
            yield p
        if i==0:
            yield pkt.getlayer("TCP")
        

    def getClassFromPkt(self, p, pkt):
        if isinstance(p, MPTCP_DSS_Map64_AckMapCsum):
            return MPTCPTest.DSS
        if isinstance(p, TCP):
            return MPTCPTest.TCPPacket

    class CapSYN(ProtoLibPacket, MPTCP_CapableSYN):
        def generate(self, s, sub=None):
            """Generate a Multipath Capable SYN segment
            
            First packet of the three-way TCP connection negotiation
            It decides of the initial sequence number and of one of the 2 keys
            used to auth the hosts in further subflows creations."""
    
            if sub is None:
                sub = s.getDefaultSubflow()

            s["snd_key"] = randintb(64)
            s["snd_token"], s["dsn"] = key2tokenAndDSN(s["snd_key"])
            sub["startseq"] = sub["seq"] = randintb(32)
            
            # generate the segment with scapy
            pkt = IP(version=4L,dst=sub["dst"], src=sub["src"])/ \
                    TCP(dport=sub["dport"],
                        sport=sub["sport"],
                        flags="S",
                        seq=sub["seq"],
                        options=[TCPOption_MP(mptcp=MPTCP_CapableSYN(
                            checksum_req=1, 
                            snd_key=s["snd_key"]))]
                        )
            sub["seq"] = (sub["seq"]+1) % (1<<32)   
    
            # Identifies this connection stage (packet type) by a name 
            s["stage"] = "MP_CAPABLE SYN"
            # The destination is expected to reply to this packet
            waitForReply = False
            return (pkt, waitForReply)
        
        def recv(self, s, pkt):
            # retrieve reply's data
            sub = s.getSubflowFromPkt(pkt)
            (l4, opt) = checkAndGetMPOption(pkt, "MP_CAPABLE") 
            
            s["rcv_key"] = opt.snd_key
            s["rcv_token"], s["data_ack"] = key2tokenAndDSN(s["rcv_key"])
            sub["rem_startseq"] = l4.seq
            sub["ack"] = l4.seq+1
    
    class CapSYNACK(ProtoLibPacket, MPTCP_CapableSYNACK): 
        def generate(self, s, sub=None):
            """Generate a Multipath Capable SYN/ACK segment
            
            Second packet of the three-way TCP connection negotiation
            """
            if sub is None:
                sub = s.getDefaultSubflow()
    
            # retrieve reply's data
            s["snd_key"] = randintb(64)
            s["snd_token"], s["dsn"] = key2tokenAndDSN(s["snd_key"])
            sub["startseq"] = sub["seq"] = randintb(32)
        
            pkt = IP(version=4L,dst=sub["dst"], src=sub["src"])/ \
                    TCP(
                        sport=sub["sport"],
                        dport=sub["dport"],
                        flags="SA",
                        seq=sub["seq"],
                        ack=sub["ack"],
                        options=[TCPOption_MP(mptcp=MPTCP_CapableSYNACK(
                            snd_key=s["snd_key"],))]
                        )
            sub["seq"] = (sub["seq"]+1) % (1<<32)   
            s["stage"] = "MP_CAPABLE SYNACK"
            waitForReply = False
            return (pkt, waitForReply)

            
        def recv(self, s, pkt):
            # retrieve reply's data
            sub = s.getSubflowFromPkt(pkt)
            (l4, opt) = checkAndGetMPOption(pkt, "MP_CAPABLE") 
        
            s["rcv_key"] = opt.snd_key
            s["rcv_token"], s["data_ack"] = key2tokenAndDSN(s["rcv_key"])
            sub["ack"] = l4.seq+1
            sub["rem_startseq"] = l4.seq # FIXME: is reached?
    
        
    class CapACK(ProtoLibPacket, MPTCP_CapableACK):
        def generate(self, s, sub=None):
            """Generate a Multipath Capable ACK segment
            
            Third packet of the three-way TCP connection negotiation
            It uses the 2nd packet (syn/ack) to generate the new sequence numbers
            (seq and ack). Also, it stores the mptcp key chosen by the other peer.
            If MPTCP cant be found in the previous packet, an exception is raised."""
            if sub is None:
                sub = s.getDefaultSubflow()
        
            pkt = IP(version=4L,dst=sub["dst"], src=sub["src"])/ \
                    TCP(
                        sport=sub["sport"],
                        dport=sub["dport"],
                        flags="A",
                        seq=sub["seq"],
                        ack=sub["ack"],
                        options=[TCPOption_MP(mptcp=MPTCP_CapableACK(
                            snd_key=s["snd_key"],
                            rcv_key=s["rcv_key"]))])
            s["stage"] = "MP_CAPABLE ACK"
            #sub["seq"] = (sub["seq"]+1) % (1<<32)   
            # No reply is expected to this packet
            waitForReply = False
            return (pkt, waitForReply)
            
        def recv(self, s, pkt):
            # retrieve reply's data
            (l4, opt) = checkAndGetMPOption(pkt, "MP_CAPABLE") 
            sub = s.getSubflowFromPkt(pkt)
            #sub["ack"] = l4.seq+1
            sub["ack"] = l4.seq
            s["rcv_token"], s["data_ack"] = key2tokenAndDSN(s["rcv_key"])
    
    
    class JoinSYN(ProtoLibPacket, MPTCP_JoinSYN):
        def generate(self, s, sub=None):
            """Generate a Multipath Join SYN segment
            
            First packet of the subflow creation negotiation
            It computes the other host's token on basis of the stored keys from
            the mptcp state. It chooses a nonce and send it along the token in the
            MP_JOIN option. Another initial data sequence number is picked as the
            subflow constitutes an other TCP connection"""
    
            if sub is None:
                sub = s.getDefaultSubflow()
    
            sub["startseq"] = sub["seq"] = randintb(32)
            sub["snd_nonce"] = randintb(32)
            sub["addr_id"] = 0
            sub["backup_flow"] = 0
        
            pkt = IP(version=4L,src=sub["src"],dst=sub["dst"])/ \
                    TCP(sport=sub["sport"],
                    dport=sub["dport"],
                    flags="S",
                    seq=sub["seq"],
                    options=[TCPOption_MP(mptcp=MPTCP_JoinSYN(
                        addr_id=sub["addr_id"],
                        backup_flow=sub["backup_flow"],
                        rcv_token=s["rcv_token"],
                        snd_nonce=sub["snd_nonce"],))])
            sub["seq"] = (sub["seq"]+1) % (1<<32)   
            s["stage"] = "MP_JOIN SYN"
            waitForReply = False
            return (pkt, waitForReply)
            
        def recv(self, s, pkt):
            (l4, opt) = checkAndGetMPOption(pkt, "MP_JOIN") 
            sub = s.getSubflowFromPkt(pkt)
            sub["ack"] = l4.seq+1
            sub["addr_id"] = opt.addr_id
            sub["backup_flow"] = opt.backup_flow
            sub["rcv_nonce"] = opt.snd_nonce
            sub["rem_startseq"] = l4.seq

    
    class JoinSYNACK(ProtoLibPacket, MPTCP_JoinSYNACK):
        def generate(self, s, sub=None):
            """Generate a Multipath Join SYN/ACK segment
    
            Second packet of the subflow creation negotiation
            In a similar way to the first packet, it picks a nonce and send it
            as a cryptographic challenge for the other host.
            It auths itself by sending a 64bits truncated MAC generated with the
            keys, the received and the newly-picked nonce.
            """
            if sub is None:
                sub = s.getDefaultSubflow()
            
            sub["startseq"] = sub["seq"] = randintb(32)
            sub["snd_nonce"] = randintb(32)
            sub["snd_mac"] = genhmac(
                                    s["snd_key"],s["rcv_key"],
                                    sub["snd_nonce"],sub["rcv_nonce"])
            
            pkt = IP(version=4L,src=sub["src"],dst=sub["dst"])/ \
                    TCP(sport=sub["sport"],
                    dport=sub["dport"],
                    flags="SA",
                    seq=sub["seq"],
                    ack=sub["ack"],
                    options=[TCPOption_MP(mptcp=MPTCP_JoinSYNACK(
                        addr_id=sub["addr_id"],
                        backup_flow=sub["backup_flow"],
                        snd_mac64=sub["snd_mac"]>>(160-64),
                        snd_nonce=sub["snd_nonce"],))
                        ])
            sub["seq"] = (sub["seq"]+1) % (1<<32)
            s["stage"] = "MP_JOIN SYNACK"
            
            waitForReply = False
            return (pkt, waitForReply)
        
        def recv(self, s, pkt):
            (l4, opt) = checkAndGetMPOption(pkt, "MP_JOIN") 
            sub = s.getSubflowFromPkt(pkt)
    
            sub["ack"] = l4.seq+1
            sub["rcv_nonce"] = opt.snd_nonce
            sub["rem_startseq"] = l4.seq

    
    class JoinACK(ProtoLibPacket, MPTCP_JoinACK):
        def generate(self, s, sub=None):
            if sub is None:
                sub = s.getDefaultSubflow()
            sub["snd_mac"] = genhmac(
                                    s["snd_key"],s["rcv_key"],
                                    sub["snd_nonce"],sub["rcv_nonce"])
            
            pkt = IP(version=4L,src=sub["src"],dst=sub["dst"])/ \
                    TCP(sport=sub["sport"],
                    dport=sub["dport"],
                    flags="A",
                    seq=sub["seq"],
                    ack=sub["ack"],
                    options=[TCPOption_MP(mptcp=MPTCP_JoinACK(
                        snd_mac=sub["snd_mac"],))])
            s["stage"] = "MP_JOIN ACK"
            #sub["seq"] = (sub["seq"]+1) % (1<<32)
            #return (pkt, False) # FIXME : should be True according to draft
            waitForReply = False
            return (pkt, waitForReply)

        def recv(self, s, pkt):
            (l4, opt) = checkAndGetMPOption(pkt, "MP_JOIN") 
            sub = s.getSubflowFromPkt(pkt)
    
            sub["ack"] = l4.seq #+1

    

    class DSSACK(ProtoLibPacket, MPTCP_DSS_Ack, MPTCP_DSS_Ack64):
        def generate(self, s, payload="", sub=None, a=False, f=False,
                waitAck=False):
            """Data Sequence Signal segment. Data are pushed in this type of
            segment. This uses the ACK-only variant"""
            if sub is None:
                sub = s.getDefaultSubflow()
            
            if a:
                data_ack = s["data_ack"]
                flags="aA"
                DSSPacket = MPTCP_DSS_Ack64 
            else:
                data_ack = get32bitSeq(s["data_ack"])
                flags = "A"
                DSSPacket = MPTCP_DSS_Ack
            if f:
                flags += "F"
            
            mptcpDSS = DSSPacket(flags=flags,
                            data_ack=data_ack)
            
            pkt = IP(version=4L,src=sub["src"],dst=sub["dst"])/ \
                    TCP(
                        sport=sub["sport"],
                        dport=sub["dport"],
                        flags="A%s"%("P" if payload else ""),
                        seq=sub["seq"],
                        ack=sub["ack"],
                        options=[TCPOption_MP(mptcp=mptcpDSS)]    
                        )/ payload

            s["stage"] = "DSS ACK"
            sub["seq"] += len(payload)
            s["dsn"] += len(payload)
            if f: s["dsn"] += 1 # draft: a DATA_FIN is accounted at data-level
            return (pkt, True if payload or f else False)

        def recv(self, s, pkt):
            (l4, opt) = checkAndGetMPOption(pkt, "DSS") 
            return MPTCPTest.TCPPacket().recv(s, pkt)


    class DSSMAP(ProtoLibPacket, MPTCP_DSS_MapCsum, MPTCP_DSS_Map64Csum): 
        def generate(self, s, payload, length, checksum, sub, m=False,
                f=False, waitAck=False):
            """Data Sequence Signal segment. Data are pushed in this type of
            segment. This uses the Map-only variant"""
            if m:
                flags="mM"
                dsn = s["dsn"]
                DSSPacket = MPTCP_DSS_Map64Csum 
            else:
                dsn = get32bitSeq(s["dsn"])
                flags = "M"
                DSSPacket = MPTCP_DSS_MapCsum
            if f:
                flags += "F"
            
            mptcpDSS = DSSPacket(flags=flags,dsn=dsn,
                            subflow_seqnum=sub["seq"]-sub["startseq"],
                            datalevel_len=length,checksum = checksum)
            
            pkt = IP(version=4L,src=sub["src"],dst=sub["dst"])/ \
                    TCP(
                        sport=sub["sport"],
                        dport=sub["dport"],
                        flags="A%s"%("P" if payload else ""),
                        seq=sub["seq"],
                        ack=sub["ack"],
                        options=[TCPOption_MP(mptcp=mptcpDSS)]    
                        )/ payload

            sub["seq"] += len(payload)
            s["dsn"] += len(payload)
            return (pkt, waitAck)

        def recv(self, s, pkt):
            (l4, opt) = checkAndGetMPOption(pkt, "DSS") 
            sub = s.getSubflowFromPkt(pkt)
            plen = len(l4.payload)
            # TCP-level FIN accounting
            if "F" in l4.sprintf("%TCP.flags%"):
                plen += 1
            sub["ack"] = l4.seq+plen
            # Data-level FIN accounting (what about combination FIN+DataFIN?)
            if flagIn(opt.flags, "F"):
                plen += 1
            if not plen>0: return
            sub["map"] = {
                    "dsn":long(long(s["data_ack"]&0xFFFFFFFF00000000)|get32bitSeq(opt.dsn)), 
                    "subseq":opt.subflow_seqnum,
                    "datalen": opt.datalevel_len}
            s["data_ack"] = getDataAckForPkt(s, sub, l4, plen,f=True if
                    flagIn(opt.flags, "F") else False)
            return MPTCPTest.DSSACK

    class DSSFIN(ProtoLibPacket):
        """Regular standalone DataFIN, a map, an data ack, no payload"""
        def generate(self, s, sub=None):
            if sub is None:
                sub = s.getDefaultSubflow()

            p = MPTCPTest.DSS()
            checksum = genDSSChecksum(s["dsn"], 0, 1, "")
            (pkt, wait) = p.generate(s, payload="", length=1, checksum=checksum, sub=sub,
                    subseq=0,
                    f=True)
            waitForReply = False
            return (pkt, waitForReply)
    DSSFINACK = DSSFIN

    
    class DSS(ProtoLibPacket, MPTCP_DSS, #MPTCP_DSS_AckMapCsum,
            MPTCP_DSS_Map64_AckMapCsum, MPTCP_DSS_Ack64MapCsum): #MPTCP_DSS_AckMap64Csum):
        def generate(self, s, payload, length, checksum, sub, subseq=-1,a=False,
                m=False, f=False, waitAck=False):
            """Data Sequence Signal segment. Data are pushed in this type of
            segment. This uses the Map+Ack variant"""
            if subseq < 0:
                subseq=sub["seq"]-sub["startseq"]
            
            flags = "AM"
            data_ack = get32bitSeq(s["data_ack"])
            dsn = get32bitSeq(s["dsn"]) 
            DSSPacket = MPTCP_DSS_AckMapCsum
            if a:
                flags+="a"
                data_ack = s["data_ack"]
                DSSPacket = MPTCP_DSS_Ack64MapCsum
            if m:
                flags+="m"
                dsn = s["dsn"]
                DSSPacket = MPTCP_DSS_AckMap64Csum 
            if a and m:
                DSSPacket = MPTCP_DSS
            if f:
                flags += "F"

            mptcpDSS = DSSPacket(flags=flags, data_ack=data_ack, dsn=dsn,
                    subflow_seqnum=subseq,
                    datalevel_len=length, checksum = checksum)
            
            pkt = IP(version=4L,src=sub["src"],dst=sub["dst"])/ \
                    TCP(
                        sport=sub["sport"],
                        dport=sub["dport"],
                        flags="A%s"%("P" if payload else ""),
                        seq=sub["seq"],
                        ack=sub["ack"],
                        options=[TCPOption_MP(mptcp=mptcpDSS)]    
                        )/ payload

            s["stage"] = "DSS MAP+ACK"
            plen = len(payload)
            sub["seq"] += plen
            s["dsn"] += plen
            if f: s["dsn"] += 1 # draft: a DATA_FIN is accounted at data-level
                                # but not at the subflow level.
            return (pkt, waitAck)
        
        def recv(self, s, pkt):
            # same as MAP since there is no handling of DATA_ACK received
            return MPTCPTest.DSSMAP().recv(s, pkt)
    
    def send_data_sub(self, s, data, sub=None, imap=[0], waitAck=False):
        """Send data using subflow sub. Split in several TCP segments if datalength is
        greater than mss"""
        if sub is None: sub = s.getDefaultSubflow()
        length = len(data)
        checksum = genDSSChecksum(s["dsn"], sub["seq"]-sub["startseq"], 
                                length, data)
        mem = [] # to retrieve the packet sent and rcvd from scenario context
        i = 0
        while data:
            payload, data = data[0:sub["mss"]], data[sub["mss"]:]
            # The map is sent along only for the (i+1)th segments from imap
            if i in imap: mem.append(self.tester.sendpkt(self.DSS, s, payload=payload,
                    length=length, checksum=checksum, sub=sub, waitAck=waitAck))
            else: mem.append(self.tester.sendpkt(self.Push, s,
                payload=payload, sub=sub, waitAck=waitAck))
            i += 1
        return mem

    
    def send_data(self, s, data="", waitAck=False):
        """Send data using all necessary subflows to minimize packets sent on each
        interface"""
        mem = [] # to retrieve the packet sent and rcvd from scenario context
        subn = len(s.sub)
        pktlen = int(math.ceil(len(data)/float(subn)))
        for subflow in s.sub:
            payload, data = data[0:pktlen], data[pktlen:]
            mem.extend(self.send_data_sub(s, data=payload, sub=subflow,
                waitAck=waitAck))
        return mem

    
    class TCPPacket(ProtoLibPacket, TCP):
        def generate(self, s, payload="", sub=None, f=False, waitAck=False):
            """Generate a regular TCP segment"""
            
            if sub is None:
                sub = s.getDefaultSubflow()
        
            flags = "A"
            if f: flags+="F"
            if payload: flags+="P"

            pkt = IP(version=4L,dst=sub["dst"], src=sub["src"])/ \
                    TCP(
                        sport=sub["sport"],
                        dport=sub["dport"],
                        flags=flags,
                        seq=sub["seq"],
                        ack=sub["ack"]) \
                    / payload
            s["stage"] = "TCP"
            plen = len(payload)
            if f: plen += 1
            sub["seq"] += plen 
            s["dsn"] += plen 
            
            #if f or waitAck: # or payload
            #    waitForReply = True
            #else:
            waitForReply = False
            return (pkt, waitForReply)
            
        def recv(self, s, pkt):
            # retrieve reply's data
            l4 = pkt.getlayer("TCP") 
            sub = s.getSubflowFromPkt(pkt)
            plen = len(l4.payload)
            if "F" in l4.sprintf("%TCP.flags%"):
                plen += 1
            sub["ack"] = l4.seq+plen
            if not plen>0: return
            s["data_ack"] = getDataAckForPkt(s, sub, l4, plen)
            return MPTCPTest.DSSACK

    ACK = TCPPacket
    
    class FINACK(ProtoLibPacket):
        """Regular standalone DataFIN, a map, an data ack, no payload"""
        def generate(self, s, sub=None):
            if sub is None:
                sub = s.getDefaultSubflow()
            p = MPTCPTest.TCPPacket()
            return p.generate(s, payload="", sub=sub, f=True)        

    class FIN(ProtoLibPacket):
        def generate(self, s, sub=None):
            """Generate a regular TCP FIN segment"""
            p = MPTCPTest.TCPPacket()
            return p.generate(s, payload="", sub=sub, f=True)
            


    class Wait(ProtoLibPacket):
        def generate(self, s, sub=None, waitfct=None, timeout=None,
                buffermode=False):
            if waitfct:
                if sub is None:
                    return (None, (waitfct, timeout, buffermode))
                return (None, (lambda pkt: s.isPacketFromSubflow(pkt) and
                        waitfct(pkt), timeout, buffermode))
            if sub is None:
                # accept packets from existing connection or new (SYN)
                return (None, (lambda pkt: pkt.haslayer(TCP) and 
                        (pkt.sprintf("%TCP.flags%") == "S" or
                        s.isPacketFromConnection(pkt)), timeout, buffermode))
            else:
                return (None, lambda pkt: sub.isPacketFromSubflow(pkt))

        @classmethod
        def filterOnSubflow(cls, sub):
            """Filter arriving packets for subflow sub.
            To put as value for waitfct argument of
            Wait.generate()"""
            return lambda pkt: sub.isPacketFromSubflow(pkt)

        @classmethod
        def filterOnConnection(cls, s):
            """Filter arriving packets for connection s.
            To put as value for waitfct argument of
            Wait.generate()"""
            return lambda pkt: s.isPacketFromConnection(pkt)

        @classmethod
        def waitAckForPkt(cls, s, oldpkt):
            """Filter arriving ACK packets if they ack oldpkt.
            (only at the TCP sequence level)"""
            sub = s.getSubflowFromPkt(oldpkt)
            def filterfct(pkt):
                l4 = pkt.getlayer(IP)
                # FIXME: no offset ?           
                if sub["seq"] == l4.ack and \
                        l4.seq == sub["ack"]:
                    return True
                return False
            return filterfct


class SubflowState(ProtoState):
    def __init__(self, mpconn, initstate={},conf=None):
        if conf: ProtoState.__init__(self, initstate=initstate, conf=conf)
        else: ProtoState.__init__(self, initstate=initstate)
        self.mpconn = mpconn # ref to mptcp connection state
        self.name = "Subflow"

    def initAttr(self):
        # per subflow
        self.d["snd_nonce"] = 0
        self.d["rcv_nonce"] = 0
        self.d["snd_mac"] = 0
        self.d["rcv_mac"] = 0
        self.d["seq"] = 0
        self.d["startseq"] = 0
        self.d["ack"] = 0
        self.d["mss"] = 25 # adjust to make more or less packets
        self.d["map"] = {}
        

    def getId(self):
        return (self.d["dst"], self.d["src"], self.d["dport"],
                self.d["sport"])

    def getIndex(self):
        return self.mpconn.sub.index(self)

    def isPacketFromSubflow(self, rawpkt):
        if not rawpkt.haslayer(TCP): return False
        pkt = rawpkt.getlayer("IP")
        tupleid = self.getId()
        self.debug("tuple: %s, should match %s"%(tupleid, (pkt.src, pkt.dst,
                pkt.sport, pkt.dport)),level=4)
        if tupleid == (pkt.dst, pkt.src, pkt.dport, pkt.sport) or \
                tupleid == (pkt.src, pkt.dst, pkt.sport, pkt.dport):
            self.debug("yes.", level=4)
            return True
        self.debug("no.", level=4)
        return False

    def invertState(self):
        """Generate a new state inverted from the current. Might be useful to
        send to other end in some cases"""
        return SubflowState(mpconn=self.mpconn, 
                initstate={"dst": self["src"], "src":self["dst"], 
                    "dport":self["sport"],"sport":self["dport"]
                })

    def __setitem__(self, attr, val):
        ProtoState.__setitem__(self, attr, val)
        if self.d["dst"] == self.d["src"]:
            raise Exception("HEY")

    def __repr__(self):
        return str(self.d)
    
    def __str__(self):
        return str(self.d)
        

class MPTCPState(ProtoState):
    def __init__(self, initstate={}, conf=None):
        if conf: ProtoState.__init__(self, initstate=initstate, conf=conf)
        else: ProtoState.__init__(self, initstate=initstate)
        self.sub = []
        self.default = 0
        self.name = "MPTCP Connection"
    
    def initAttr(self):
        # auth
        self.d["rcv_token"] = 0
        self.d["snd_token"] = 0
        self.d["snd_key"] = 0
        self.d["rcv_key"] = 0
        
        # DSS related
        self.d["dsn"] = 0
        self.d["subflow_seqnum"] = 0
        self.d["datalevel_len"] = 0
        self.d["data_ack"] = 0

    def createSubflow(self, dst, src, dport=80, sport=0):
        if sport == 0: sport = random.randrange(1025,2<<15)
        return SubflowState(mpconn=self, 
                initstate={"dst":dst, "src":src, "dport":dport,"sport":sport})

    def registerSubflow(self, ss):
        for sub in self.sub:
            if ss.getId() == sub.getId():
                self.debug("subflow %s already exists" % ss.getId())
                return sub
        else:
            self.sub.append(ss)
            ss.name = "Subflow #%i" % ss.getIndex()
        return ss

    def registerNewSubflow(self, **kargs):
        return self.registerSubflow(self.createSubflow(**kargs))

    def getSubflow(self, subflow):
        if type(subflow) is int:
            if subflow < len(self.sub):
                return self.sub[subflow]
        return None

    def getDefaultSubflow(self):
        return self.sub[self.default]

    def setDefaultSubflow(self, ss):
        self.default = self.sub.index(ss)

    def newSubflowFromRcvdPkt(self, pkt):
        """create a sufblow state from data received in previous packet.
        and return its pair (id, subflowstate)"""
        return self.registerNewSubflow(dst=pkt.src, src=pkt.dst, dport=pkt.sport,
                sport=pkt.dport)

    def getSubflowFromPkt(self, pkt):
        """return the subflow on which the packet has been received"""
        if self.sub:
            for sub in self.sub:
                if sub.isPacketFromSubflow(pkt):
                    return sub
        return self.newSubflowFromRcvdPkt(pkt)
    
    def isPacketFromConnection(self, pkt):
        for sub in self.sub:
            if sub.isPacketFromSubflow(pkt):
                return True
        return False

    def invertState(self):
        """Generate a new state inverted from the current. Might be useful to
        send to other end in some cases"""
        return MPTCPState({
                "snd_key":self["rcv_key"], "rcv_key": self["snd_key"], 
                "dst": self["src"], "src":self["dst"], "dport":self["dport"],
                })
    
    def update(self, extrastate):
        """Update the current state with the extrastate. Extrastate must be a
        ProtoState derivative"""
        ProtoState.update(self, extrastate)
        if type(extrastate) is type(self):
            self.sub = extrastate.sub
        return self
    
    def logPacket(self, pkt):
        self.d["prev_pkt"] = pkt
        sub = self.getSubflowFromPkt(pkt)
        self.setDefaultSubflow(sub) # impose to reply to last-active subflow
        # when subflow to send on is not explicit
        sub.logPacket(pkt)


# vim: set ts=4 sts=4 sw=4 et:
