# -*- coding: iso-8859-1 -*-
""" crypto.cipher.tkip_encr

    TKIP encryption from IEEE 802.11 TGi

    TKIP uses WEP (ARC4 with crc32) and key mixing
    This is only the encryption and not the Michael integrity check!

    Copyright © (c) 2002 by Paul A. Lambert
    Read LICENSE.txt for license information.

    November 2002
"""
from crypto.cipher.arc4 import ARC4
from zlib   import crc32
from struct import pack
from crypto.keyedHash.tkip_key_mixing import TKIP_Mixer
from crypto.errors import BadKeySizeError, IntegrityCheckError
from binascii_plus import *

class TKIP_encr:
    """ TKIP Stream Cipher Algorithm without the Michael integrity check

        TKIP encryption on an MPDU using WEP with a longer 'iv'
        and the TKIP key mixing algorithm .  This does NOT include
        the Michael integrity algorithm that operates on the MSDU data.
    """
    def __init__(self, key=None, transmitterAddress=None, keyID=None):
        """ Initialize TKIP_encr, key -> octet string for key """
        assert(keyID==0 or keyID==None), 'keyID should be zero in TKIP'
        self.keyId = 0
        self.name  = 'TKIP_encr'
        self.strength    = 128
        self.encryptHeaderSize = 8        # used to skip octets on decrypt
        self.arc4 = ARC4()                # base algorithm
        self.keyMixer = TKIP_Mixer(key,transmitterAddress)
        if key != None:                   # normally in base init, uses adjusted keySize
            self.setKey(key)
        if transmitterAddress!=None:
            self.setTA(transmitterAddress)

    def setKey(self, key, ta=None):
        """ Set key, key string is 16 octets long """
        self.keyMixer.setKey(key)
        if ta!=None:
            self.setTA(ta)

    def setTA(self, transmitterAddress):
        """ Set the transmitter address """
        self.keyMixer.setTA(transmitterAddress)

    def _getIVandKeyID(self, cipherText):
        """ Parse the TKIP header to get iv and set KeyID
            iv is returned as octet string and is little-endian!!!
        """
        assert(ord(cipherText[3])& 0x20),'extIV SHOULD be set in TKIP header'
        self.setCurrentKeyID = (ord(cipherText[3])&0xC0)>>6
        return cipherText[:3] + cipherText[5:9] # note iv octets are little-endian!!!

    def _makeARC4key(self, tscOctets, keyID=0):
        """ Make an ARC4 key from TKIP Sequence Counter Octets (little-endian) """
        if keyID!=0 :
            raise 'TKIP expects keyID of zero'
        print "tscOctets in tkmixer=",b2a_p(tscOctets)
        newKey = self.keyMixer.newKey(tscOctets)
        print "newKey=", b2a_p(newKey)
        return newKey

    def encrypt(self, plainText, iv):
        """ Encrypt a string and return a binary string
            iv is 6 octets of little-endian encoded pn
        """
        assert(len(iv)==6),'TKIP bad IV size on encryption'
        self.pnField = iv
        self.arc4.setKey( self._makeARC4key(iv) )
        eh1 = chr((ord(iv[0])|0x20)&0x7f)
        encryptionHeader = iv[0] + eh1 + iv[1] + chr((self.keyId<<6)|0x20) + iv[2:]
        crc = pack('<I', crc32(plainText) )
        cipherText = encryptionHeader + self.arc4.encrypt(plainText+crc)
        return cipherText

    def decrypt(self, cipherText):
        """ Decrypt a WEP packet, assumes WEP 4 byte header on packet """
        assert(ord(cipherText[3])& 0x20),'extIV SHOULD be set in TKIP header'
        self.setCurrentKeyID = (ord(cipherText[3])&0xC0)>>6
        iv = cipherText[0]+cipherText[2]+cipherText[4:8]
        self.pnField = iv
        self.arc4.setKey( self._makeARC4key(iv) )
        plainText = self.arc4.decrypt(cipherText[self.encryptHeaderSize:])
        if plainText[-4:] != pack('<I',crc32(plainText[:-4])):  # check data integrity
            raise IntegrityCheckError, 'WEP CRC Integrity Check Error'
        return plainText[:-4]