Add boundary check for 802.11 decryption
[metze/wireshark/wip.git] / epan / crypt / airpdcap.c
index 757ab18f82eac4f64b2b3525ffacb88fb09806dd..e584bf43acce3923c4cdfe30edffd0c1526a76e2 100644 (file)
@@ -1,6 +1,5 @@
 /* airpdcap.c
  *
- * $Id$
  * Copyright (c) 2006 CACE Technologies, Davis (California)
  * All rights reserved.
  *
  *    may be used to endorse or promote products derived from this software
  *    without specific prior written permission.
  *
+ * Alternatively, this software may be distributed under the terms of the
+ * GNU General Public License ("GPL") version 2 as published by the Free
+ * Software Foundation.
+ *
  * THIS SOFTWARE IS PROVIDED BY THE PROJECT AND CONTRIBUTORS ``AS IS'' AND
  * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  * SUCH DAMAGE.
  */
 
+/*
+ * The files matching airpcap*.[ch] were originally developed as part of
+ * Wireshark's support for AirPcap adapters. However, they've been used
+ * for general 802.11 decryption for quite some time. It might make sense
+ * to rename them accordingly.
+ */
+
 /****************************************************************************/
-/*     File includes                                                           */
+/*      File includes                                                       */
+
+#include "config.h"
+
+#include <glib.h>
+
+#include <wsutil/crc32.h>
+#include <wsutil/rc4.h>
+#include <wsutil/sha1.h>
+#include <wsutil/md5.h>
+#include <wsutil/pint.h>
 
 #include <epan/tvbuff.h>
-#include <epan/crc32.h>
+#include <epan/to_str.h>
 #include <epan/strutil.h>
-#include <epan/ws_strsplit.h>
-#include <epan/emem.h>
-#include <epan/pint.h>
+#include <epan/crypt/airpdcap_rijndael.h>
 
 #include "airpdcap_system.h"
 #include "airpdcap_int.h"
 
-#include "crypt-sha1.h"
-#include "crypt-md5.h"
-
 #include "airpdcap_debug.h"
 
 #include "wep-wpadefs.h"
 
-/****************************************************************************/
 
 /****************************************************************************/
-/*     Constant definitions                                                    */
 
-#define        AIRPDCAP_SHA_DIGEST_LEN 20
+/****************************************************************************/
+/*      Constant definitions                                                    */
 
-/*     EAPOL definitions                                                       */
+/*      EAPOL definitions                                                       */
 /**
  * Length of the EAPOL-Key key confirmation key (KCK) used to calculate
  * MIC over EAPOL frame and validate an EAPOL packet (128 bits)
  */
-#define        AIRPDCAP_WPA_KCK_LEN    16
+#define AIRPDCAP_WPA_KCK_LEN    16
 /**
  *Offset of the Key MIC in the EAPOL packet body
  */
-#define        AIRPDCAP_WPA_MICKEY_OFFSET      77
+#define AIRPDCAP_WPA_MICKEY_OFFSET      77
 /**
  * Maximum length of the EAPOL packet (it depends on the maximum MAC
  * frame size)
  */
-#define        AIRPDCAP_WPA_MAX_EAPOL_LEN      4095
+#define AIRPDCAP_WPA_MAX_EAPOL_LEN      4095
 /**
  * EAPOL Key Descriptor Version 1, used for all EAPOL-Key frames to and
  * from a STA when neither the group nor pairwise ciphers are CCMP for
@@ -78,7 +92,7 @@
  * @note
  * Defined in 802.11i-2004, page 78
  */
-#define        AIRPDCAP_WPA_KEY_VER_CCMP       1
+#define AIRPDCAP_WPA_KEY_VER_NOT_CCMP   1
 /**
  * EAPOL Key Descriptor Version 2, used for all EAPOL-Key frames to and
  * from a STA when either the pairwise or the group cipher is AES-CCMP
  * /note
  * Defined in 802.11i-2004, page 78
  */
-#define        AIRPDCAP_WPA_KEY_VER_AES_CCMP   2
+#define AIRPDCAP_WPA_KEY_VER_AES_CCMP   2
+
+/** Define EAPOL Key Descriptor type values:  use 254 for WPA and 2 for WPA2 **/
+#define AIRPDCAP_RSN_WPA_KEY_DESCRIPTOR 254
+#define AIRPDCAP_RSN_WPA2_KEY_DESCRIPTOR 2
 
 /****************************************************************************/
 
+
+
 /****************************************************************************/
-/*     Macro definitions                                                       */
+/*      Macro definitions                                                       */
 
 extern const UINT32 crc32_table[256];
-#define CRC(crc, ch)    (crc = (crc >> 8) ^ crc32_table[(crc ^ (ch)) & 0xff])
+#define CRC(crc, ch)     (crc = (crc >> 8) ^ crc32_table[(crc ^ (ch)) & 0xff])
 
-#define        AIRPDCAP_GET_TK(ptk)    (ptk + 32)
+#define AIRPDCAP_GET_TK(ptk)    (ptk + 32)
 
 /****************************************************************************/
 
 /****************************************************************************/
-/*     Type definitions                                                        */
+/*      Type definitions                                                        */
 
-/*     Internal function prototype declarations                                */
+/*      Internal function prototype declarations                                */
 
-#ifdef __cplusplus
+#ifdef  __cplusplus
 extern "C" {
 #endif
 
 /**
  * It is a step of the PBKDF2 (specifically the PKCS #5 v2.0) defined in
  * the RFC 2898 to derive a key (used as PMK in WPA)
- * @param password [IN] pointer to a password (sequence of between 8 and
+ * @param ppbytes [IN] pointer to a password (sequence of between 8 and
  * 63 ASCII encoded characters)
  * @param ssid [IN] pointer to the SSID string encoded in max 32 ASCII
  * encoded characters
  * @param iterations [IN] times to hash the password (4096 for WPA)
  * @param count [IN] ???
  * @param output [OUT] pointer to a preallocated buffer of
- * AIRPDCAP_SHA_DIGEST_LEN characters that will contain a part of the key
+ * SHA1_DIGEST_LEN characters that will contain a part of the key
  */
 static INT AirPDcapRsnaPwd2PskStep(
     const guint8 *ppbytes,
@@ -135,7 +155,7 @@ static INT AirPDcapRsnaPwd2PskStep(
  * It calculates the passphrase-to-PSK mapping reccomanded for use with
  * RSNAs. This implementation uses the PBKDF2 method defined in the RFC
  * 2898.
- * @param password [IN] pointer to a password (sequence of between 8 and
+ * @param passphrase [IN] pointer to a password (sequence of between 8 and
  * 63 ASCII encoded characters)
  * @param ssid [IN] pointer to the SSID string encoded in max 32 ASCII
  * encoded characters
@@ -173,7 +193,6 @@ static INT AirPDcapRsna4WHandshake(
     PAIRPDCAP_CONTEXT ctx,
     const UCHAR *data,
     AIRPDCAP_SEC_ASSOCIATION *sa,
-    PAIRPDCAP_KEY_ITEM key,
     INT offset)
     ;
 /**
@@ -227,40 +246,249 @@ static void AirPDcapRsnaPrfX(
     AIRPDCAP_SEC_ASSOCIATION *sa,
     const UCHAR pmk[32],
     const UCHAR snonce[32],
-    const INT x,       /*      for TKIP 512, for CCMP 384      */
+    const INT x,        /*      for TKIP 512, for CCMP 384      */
     UCHAR *ptk)
     ;
 
-#ifdef __cplusplus
+#ifdef  __cplusplus
 }
 #endif
 
 /****************************************************************************/
 
 /****************************************************************************/
-/* Exported function definitions                                               */
+/* Exported function definitions                                                */
 
-#ifdef __cplusplus
+#ifdef  __cplusplus
 extern "C" {
 #endif
 
-INT AirPDcapPacketProcess(
+const guint8 broadcast_mac[] = { 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF };
+
+
+/* NOTE : this assumes the WPA RSN IE format.  If it were to be a generic RSN IE, then
+   we would need to change the structure since it could be variable length depending on the number
+   of unicast OUI and auth OUI. */
+typedef struct {
+    guint8 bElementID;
+    guint8 bLength;
+    guint8  OUI[4];
+    guint16 iVersion;
+    guint8  multicastOUI[4];
+    guint16 iUnicastCount;      /* this should always be 1 for WPA client */
+    guint8  unicastOUI[4];
+    guint16 iAuthCount;         /* this should always be 1 for WPA client */
+    guint8  authOUI[4];
+    guint16 iWPAcap;
+} RSN_IE;
+
+#define EAPKEY_MIC_LEN  16  /* length of the MIC key for EAPoL_Key packet's MIC using MD5 */
+#define NONCE_LEN 32
+
+#define TKIP_GROUP_KEY_LEN 32
+#define CCMP_GROUP_KEY_LEN 16
+/* Minimum size of the key bytes payload for a TKIP group key in an M3 message*/
+#define TKIP_GROUP_KEYBYTES_LEN ( sizeof(RSN_IE) + 8 + TKIP_GROUP_KEY_LEN + 6 ) /* 72 */
+/* arbitrary upper limit */
+#define TKIP_GROUP_KEYBYTES_LEN_MAX ( TKIP_GROUP_KEYBYTES_LEN + 28 )
+/* Minimum size of the key bytes payload for a TKIP group key in a group key message */
+#define TKIP_GROUP_KEYBYTES_LEN_GKEY (8 + 8 + TKIP_GROUP_KEY_LEN ) /* 48 */
+/*  size of CCMP key bytes payload */
+#define CCMP_GROUP_KEYBYTES_LEN ( sizeof(RSN_IE) + 8 + CCMP_GROUP_KEY_LEN + 6 ) /* 56 */
+typedef struct {
+    guint8  type;
+    guint8  key_information[2];  /* Make this an array to avoid alignment issues */
+    guint8  key_length[2];  /* Make this an array to avoid alignment issues */
+    guint8  replay_counter[8];
+    guint8  key_nonce[NONCE_LEN];
+    guint8  key_iv[16];
+    guint8  key_sequence_counter[8];  /* also called the RSC */
+    guint8  key_id[8];
+    guint8  key_mic[EAPKEY_MIC_LEN];
+    guint8  key_data_len[2];  /* Make this an array rather than a U16 to avoid alignment shifting */
+    guint8  ie[TKIP_GROUP_KEYBYTES_LEN_MAX]; /* Make this an array to avoid alignment issues */
+} EAPOL_RSN_KEY,  * P_EAPOL_RSN_KEY;
+#define RSN_KEY_WITHOUT_KEYBYTES_LEN sizeof(EAPOL_RSN_KEY)-TKIP_GROUP_KEYBYTES_LEN_MAX
+/* Minimum possible group key msg size (group key msg using CCMP as cipher)*/
+#define GROUP_KEY_PAYLOAD_LEN_MIN RSN_KEY_WITHOUT_KEYBYTES_LEN+CCMP_GROUP_KEY_LEN
+
+/* XXX - what if this doesn't get the key? */
+static INT
+AirPDcapDecryptWPABroadcastKey(const EAPOL_RSN_KEY *pEAPKey, guint8  *decryption_key, PAIRPDCAP_SEC_ASSOCIATION sa, gboolean group_hshake)
+{
+    guint8 key_version;
+    guint8  *szEncryptedKey;
+    guint16 key_bytes_len = 0; /* Length of the total key data field */
+    guint16 key_len;           /* Actual group key length */
+    static AIRPDCAP_KEY_ITEM dummy_key; /* needed in case AirPDcapRsnaMng() wants the key structure */
+    AIRPDCAP_SEC_ASSOCIATION *tmp_sa;
+
+    /* We skip verifying the MIC of the key. If we were implementing a WPA supplicant we'd want to verify, but for a sniffer it's not needed. */
+
+    /* Preparation for decrypting the group key -  determine group key data length */
+    /* depending on whether the pairwise key is TKIP or AES encryption key */
+    key_version = AIRPDCAP_EAP_KEY_DESCR_VER(pEAPKey->key_information[1]);
+    if (key_version == AIRPDCAP_WPA_KEY_VER_NOT_CCMP){
+        /* TKIP */
+        key_bytes_len = pntoh16(pEAPKey->key_length);
+    }else if (key_version == AIRPDCAP_WPA_KEY_VER_AES_CCMP){
+        /* AES */
+        key_bytes_len = pntoh16(pEAPKey->key_data_len);
+
+        /* AES keys must be at least 128 bits = 16 bytes. */
+        if (key_bytes_len < 16) {
+            return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
+        }
+    }
+
+    if (key_bytes_len > TKIP_GROUP_KEYBYTES_LEN_MAX || key_bytes_len == 0) { /* Don't read past the end of pEAPKey->ie */
+        return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
+    }
+
+    /* Encrypted key is in the information element field of the EAPOL key packet */
+    szEncryptedKey = (guint8 *)g_memdup(pEAPKey->ie, key_bytes_len);
+
+    DEBUG_DUMP("Encrypted Broadcast key:", szEncryptedKey, key_bytes_len);
+    DEBUG_DUMP("KeyIV:", pEAPKey->key_iv, 16);
+    DEBUG_DUMP("decryption_key:", decryption_key, 16);
+
+    /* We are rekeying, save old sa */
+    tmp_sa=(AIRPDCAP_SEC_ASSOCIATION *)g_malloc(sizeof(AIRPDCAP_SEC_ASSOCIATION));
+    memcpy(tmp_sa, sa, sizeof(AIRPDCAP_SEC_ASSOCIATION));
+    sa->next=tmp_sa;
+
+    /* As we have no concept of the prior association request at this point, we need to deduce the     */
+    /* group key cipher from the length of the key bytes. In WPA this is straightforward as the        */
+    /* keybytes just contain the GTK, and the GTK is only in the group handshake, NOT the M3.          */
+    /* In WPA2 its a little more tricky as the M3 keybytes contain an RSN_IE, but the group handshake  */
+    /* does not. Also there are other (variable length) items in the keybytes which we need to account */
+    /* for to determine the true key length, and thus the group cipher.                                */
+
+    if (key_version == AIRPDCAP_WPA_KEY_VER_NOT_CCMP){
+        guint8 new_key[32];
+        guint8 dummy[256];
+        /* TKIP key */
+        /* Per 802.11i, Draft 3.0 spec, section 8.5.2, p. 97, line 4-8, */
+        /* group key is decrypted using RC4.  Concatenate the IV with the 16 byte EK (PTK+16) to get the decryption key */
+
+        rc4_state_struct rc4_state;
+
+        /* The WPA group key just contains the GTK bytes so deducing the type is straightforward   */
+        /* Note - WPA M3 doesn't contain a group key so we'll only be here for the group handshake */
+        sa->wpa.key_ver = (key_bytes_len >=TKIP_GROUP_KEY_LEN)?AIRPDCAP_WPA_KEY_VER_NOT_CCMP:AIRPDCAP_WPA_KEY_VER_AES_CCMP;
+
+        /* Build the full decryption key based on the IV and part of the pairwise key */
+        memcpy(new_key, pEAPKey->key_iv, 16);
+        memcpy(new_key+16, decryption_key, 16);
+        DEBUG_DUMP("FullDecrKey:", new_key, 32);
+
+        crypt_rc4_init(&rc4_state, new_key, sizeof(new_key));
+
+        /* Do dummy 256 iterations of the RC4 algorithm (per 802.11i, Draft 3.0, p. 97 line 6) */
+        crypt_rc4(&rc4_state, dummy, 256);
+        crypt_rc4(&rc4_state, szEncryptedKey, key_bytes_len);
+
+    } else if (key_version == AIRPDCAP_WPA_KEY_VER_AES_CCMP){
+        /* AES CCMP key */
+
+        guint8 key_found;
+        guint16 key_index;
+        guint8 *decrypted_data;
+
+        /* If this EAPOL frame is part of a separate group key handshake then this contains no    */
+        /* RSN IE, so we can deduct that from the calculation.                                    */
+        if (group_hshake)
+            sa->wpa.key_ver = (key_bytes_len >= (TKIP_GROUP_KEYBYTES_LEN_GKEY))?AIRPDCAP_WPA_KEY_VER_NOT_CCMP:AIRPDCAP_WPA_KEY_VER_AES_CCMP;
+        else
+            sa->wpa.key_ver = (key_bytes_len >= (TKIP_GROUP_KEYBYTES_LEN))?AIRPDCAP_WPA_KEY_VER_NOT_CCMP:AIRPDCAP_WPA_KEY_VER_AES_CCMP;
+
+        /* Unwrap the key; the result is key_bytes_len in length */
+        decrypted_data = AES_unwrap(decryption_key, 16, szEncryptedKey,  key_bytes_len);
+
+        /* With WPA2 what we get after Broadcast Key decryption is an actual RSN structure.
+           The key itself is stored as a GTK KDE
+           WPA2 IE (1 byte) id = 0xdd, length (1 byte), GTK OUI (4 bytes), key index (1 byte) and 1 reserved byte. Thus we have to
+           pass pointer to the actual key with 8 bytes offset */
+
+        key_found = FALSE;
+        key_index = 0;
+        while(key_index < key_bytes_len && !key_found){
+            guint8 rsn_id;
+
+            /* Get RSN ID */
+            rsn_id = decrypted_data[key_index];
+
+            if (rsn_id != 0xdd){
+                if (key_index+1 >= key_bytes_len){
+                    return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
+                }
+                key_index += decrypted_data[key_index+1]+2;
+            }else{
+                key_found = TRUE;
+            }
+        }
+
+        if (key_found){
+            if (key_index+8 >= key_bytes_len)
+                return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
+
+            /* Skip over the GTK header info, and don't copy past the end of the encrypted data */
+            memcpy(szEncryptedKey, decrypted_data+key_index+8, key_bytes_len-key_index-8);
+        }
+
+        g_free(decrypted_data);
+    }
+
+    key_len = (sa->wpa.key_ver==AIRPDCAP_WPA_KEY_VER_NOT_CCMP)?TKIP_GROUP_KEY_LEN:CCMP_GROUP_KEY_LEN;
+
+    /* Decrypted key is now in szEncryptedKey with len of key_len */
+    DEBUG_DUMP("Broadcast key:", szEncryptedKey, key_len);
+
+    /* Load the proper key material info into the SA */
+    sa->key = &dummy_key;  /* we just need key to be not null because it is checked in AirPDcapRsnaMng().  The WPA key materials are actually in the .wpa structure */
+    sa->validKey = TRUE;
+
+    /* Since this is a GTK and its size is only 32 bytes (vs. the 64 byte size of a PTK), we fake it and put it in at a 32-byte offset so the  */
+    /* AirPDcapRsnaMng() function will extract the right piece of the GTK for decryption. (The first 16 bytes of the GTK are used for decryption.) */
+    memset(sa->wpa.ptk, 0, sizeof(sa->wpa.ptk));
+    memcpy(sa->wpa.ptk+32, szEncryptedKey, key_len);
+    g_free(szEncryptedKey);
+    return AIRPDCAP_RET_SUCCESS_HANDSHAKE;
+}
+
+
+/* Return a pointer the the requested SA. If it doesn't exist create it. */
+static PAIRPDCAP_SEC_ASSOCIATION
+AirPDcapGetSaPtr(
+    PAIRPDCAP_CONTEXT ctx,
+    AIRPDCAP_SEC_ASSOCIATION_ID *id)
+{
+    int sa_index;
+
+    /* search for a cached Security Association for supplied BSSID and STA MAC  */
+    if ((sa_index=AirPDcapGetSa(ctx, id))==-1) {
+        /* create a new Security Association if it doesn't currently exist      */
+        if ((sa_index=AirPDcapStoreSa(ctx, id))==-1) {
+            return NULL;
+        }
+    }
+    /* get the Security Association structure   */
+    return &ctx->sa[sa_index];
+}
+
+static INT AirPDcapScanForKeys(
     PAIRPDCAP_CONTEXT ctx,
     const guint8 *data,
     const guint mac_header_len,
     const guint tot_len,
-    UCHAR *decrypt_data,
-    guint *decrypt_len,
-    PAIRPDCAP_KEY_ITEM key,
-    gboolean mngHandshake,
-    gboolean mngDecrypt)
+    AIRPDCAP_SEC_ASSOCIATION_ID id
+)
 {
-    const UCHAR *address;
-    AIRPDCAP_SEC_ASSOCIATION_ID id;
-    int index;
+    const UCHAR *addr;
+    guint bodyLength;
+    PAIRPDCAP_SEC_ASSOCIATION sta_sa;
     PAIRPDCAP_SEC_ASSOCIATION sa;
     int offset = 0;
-    guint bodyLength;
     const guint8 dot1x_header[] = {
         0xAA,             /* DSAP=SNAP */
         0xAA,             /* SSAP=SNAP */
@@ -268,9 +496,146 @@ INT AirPDcapPacketProcess(
         0x00, 0x00, 0x00, /* Org. code=encaps. Ethernet */
         0x88, 0x8E        /* Type: 802.1X authentication */
     };
+    const guint8 bt_dot1x_header[] = {
+        0xAA,             /* DSAP=SNAP */
+        0xAA,             /* SSAP=SNAP */
+        0x03,             /* Control field=Unnumbered frame */
+        0x00, 0x19, 0x58, /* Org. code=Bluetooth SIG */
+        0x00, 0x03        /* Type: Bluetooth Security */
+    };
+
+    const EAPOL_RSN_KEY *pEAPKey;
+#ifdef _DEBUG
+#define MSGBUF_LEN 255
+    CHAR msgbuf[MSGBUF_LEN];
+#endif
+    AIRPDCAP_DEBUG_TRACE_START("AirPDcapScanForKeys");
+
+    /* cache offset in the packet data */
+    offset = mac_header_len;
+
+    /* check if the packet has an LLC header and the packet is 802.1X authentication (IEEE 802.1X-2004, pg. 24) */
+    if (memcmp(data+offset, dot1x_header, 8) == 0 || memcmp(data+offset, bt_dot1x_header, 8) == 0) {
+
+        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapScanForKeys", "Authentication: EAPOL packet", AIRPDCAP_DEBUG_LEVEL_3);
+
+        /* skip LLC header */
+        offset+=8;
+
+        /* check if the packet is a EAPOL-Key (0x03) (IEEE 802.1X-2004, pg. 25) */
+        if (data[offset+1]!=3) {
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapScanForKeys", "Not EAPOL-Key", AIRPDCAP_DEBUG_LEVEL_3);
+            return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
+        }
+
+        /* get and check the body length (IEEE 802.1X-2004, pg. 25) */
+        bodyLength=pntoh16(data+offset+2);
+        if ((tot_len-offset-4) < bodyLength) { /* Only check if frame is long enough for eapol header, ignore tailing garbage, see bug 9065 */
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapScanForKeys", "EAPOL body too short", AIRPDCAP_DEBUG_LEVEL_3);
+            return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
+        }
+
+        /* skip EAPOL MPDU and go to the first byte of the body */
+        offset+=4;
+
+        pEAPKey = (const EAPOL_RSN_KEY *) (data+offset);
+
+        /* check if the key descriptor type is valid (IEEE 802.1X-2004, pg. 27) */
+        if (/*pEAPKey->type!=0x1 &&*/ /* RC4 Key Descriptor Type (deprecated) */
+            pEAPKey->type != AIRPDCAP_RSN_WPA2_KEY_DESCRIPTOR &&             /* IEEE 802.11 Key Descriptor Type  (WPA2) */
+            pEAPKey->type != AIRPDCAP_RSN_WPA_KEY_DESCRIPTOR)           /* 254 = RSN_KEY_DESCRIPTOR - WPA,              */
+        {
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapScanForKeys", "Not valid key descriptor type", AIRPDCAP_DEBUG_LEVEL_3);
+            return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
+        }
+
+        /* start with descriptor body */
+        offset+=1;
+
+        /* search for a cached Security Association for current BSSID and AP */
+        sa = AirPDcapGetSaPtr(ctx, &id);
+        if (sa == NULL){
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapScanForKeys", "No SA for BSSID found", AIRPDCAP_DEBUG_LEVEL_3);
+            return AIRPDCAP_RET_REQ_DATA;
+        }
+
+        /* It could be a Pairwise Key exchange, check */
+        if (AirPDcapRsna4WHandshake(ctx, data, sa, offset) == AIRPDCAP_RET_SUCCESS_HANDSHAKE)
+            return AIRPDCAP_RET_SUCCESS_HANDSHAKE;
+
+        if (mac_header_len + GROUP_KEY_PAYLOAD_LEN_MIN > tot_len) {
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapScanForKeys", "Message too short for Group Key", AIRPDCAP_DEBUG_LEVEL_3);
+            return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
+        }
+
+        /* Verify the bitfields: Key = 0(groupwise) Mic = 1 Ack = 1 Secure = 1 */
+        if (AIRPDCAP_EAP_KEY(data[offset+1])!=0 ||
+            AIRPDCAP_EAP_ACK(data[offset+1])!=1 ||
+            AIRPDCAP_EAP_MIC(data[offset]) != 1 ||
+            AIRPDCAP_EAP_SEC(data[offset]) != 1){
+
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapScanForKeys", "Key bitfields not correct for Group Key", AIRPDCAP_DEBUG_LEVEL_3);
+            return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
+        }
+
+        /* force STA address to be the broadcast MAC so we create an SA for the groupkey */
+        memcpy(id.sta, broadcast_mac, AIRPDCAP_MAC_LEN);
+
+        /* get the Security Association structure for the broadcast MAC and AP */
+        sa = AirPDcapGetSaPtr(ctx, &id);
+        if (sa == NULL){
+            return AIRPDCAP_RET_REQ_DATA;
+        }
+
+        /* Get the SA for the STA, since we need its pairwise key to decrpyt the group key */
+
+        /* get STA address */
+        if ( (addr=AirPDcapGetStaAddress((const AIRPDCAP_MAC_FRAME_ADDR4 *)(data))) != NULL) {
+            memcpy(id.sta, addr, AIRPDCAP_MAC_LEN);
+#ifdef _DEBUG
+            g_snprintf(msgbuf, MSGBUF_LEN, "ST_MAC: %2X.%2X.%2X.%2X.%2X.%2X\t", id.sta[0],id.sta[1],id.sta[2],id.sta[3],id.sta[4],id.sta[5]);
+#endif
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapScanForKeys", msgbuf, AIRPDCAP_DEBUG_LEVEL_3);
+        } else {
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapScanForKeys", "SA not found", AIRPDCAP_DEBUG_LEVEL_5);
+            return AIRPDCAP_RET_REQ_DATA;
+        }
+
+        sta_sa = AirPDcapGetSaPtr(ctx, &id);
+        if (sta_sa == NULL){
+            return AIRPDCAP_RET_REQ_DATA;
+        }
+
+        /* Try to extract the group key and install it in the SA */
+        return (AirPDcapDecryptWPABroadcastKey(pEAPKey, sta_sa->wpa.ptk+16, sa, TRUE));
+
+    }else{
+        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapScanForKeys", "Skipping: not an EAPOL packet", AIRPDCAP_DEBUG_LEVEL_3);
+    }
+
+    AIRPDCAP_DEBUG_TRACE_END("AirPDcapScanForKeys");
+    return AIRPDCAP_RET_NO_VALID_HANDSHAKE;;
+}
+
+
+INT AirPDcapPacketProcess(
+    PAIRPDCAP_CONTEXT ctx,
+    const guint8 *data,
+    const guint mac_header_len,
+    const guint tot_len,
+    UCHAR *decrypt_data,
+    guint *decrypt_len,
+    PAIRPDCAP_KEY_ITEM key,
+    gboolean scanHandshake)
+{
+    const UCHAR *addr;
+    AIRPDCAP_SEC_ASSOCIATION_ID id;
+    UCHAR tmp_data[AIRPDCAP_MAX_CAPLEN];
+    guint tmp_len;
 
 #ifdef _DEBUG
-    CHAR msgbuf[255];
+#define MSGBUF_LEN 255
+    CHAR msgbuf[MSGBUF_LEN];
 #endif
 
     AIRPDCAP_DEBUG_TRACE_START("AirPDcapPacketProcess");
@@ -278,148 +643,129 @@ INT AirPDcapPacketProcess(
     if (ctx==NULL) {
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "NULL context", AIRPDCAP_DEBUG_LEVEL_5);
         AIRPDCAP_DEBUG_TRACE_END("AirPDcapPacketProcess");
-        return AIRPDCAP_RET_UNSUCCESS;
+        return AIRPDCAP_RET_REQ_DATA;
     }
     if (data==NULL || tot_len==0) {
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "NULL data or length=0", AIRPDCAP_DEBUG_LEVEL_5);
         AIRPDCAP_DEBUG_TRACE_END("AirPDcapPacketProcess");
-        return AIRPDCAP_RET_UNSUCCESS;
+        return AIRPDCAP_RET_REQ_DATA;
     }
 
-    /* check if the packet is of data type     */
+    /* check if the packet is of data type */
     if (AIRPDCAP_TYPE(data[0])!=AIRPDCAP_TYPE_DATA) {
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "not data packet", AIRPDCAP_DEBUG_LEVEL_5);
         return AIRPDCAP_RET_NO_DATA;
     }
 
-    /* check correct packet size, to avoid wrong elaboration of encryption algorithms  */
+    /* check correct packet size, to avoid wrong elaboration of encryption algorithms */
     if (tot_len < (UINT)(mac_header_len+AIRPDCAP_CRYPTED_DATA_MINLEN)) {
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "minimum length violated", AIRPDCAP_DEBUG_LEVEL_5);
         return AIRPDCAP_RET_WRONG_DATA_SIZE;
     }
 
+    /* Assume that the decrypt_data field is at least this size. */
+    if (tot_len > AIRPDCAP_MAX_CAPLEN) {
+        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "length too large", AIRPDCAP_DEBUG_LEVEL_3);
+        return AIRPDCAP_RET_UNSUCCESS;
+    }
+
     /* get BSSID */
-    if ( (address=AirPDcapGetBssidAddress((const AIRPDCAP_MAC_FRAME_ADDR4 *)(data))) != NULL) {
-        memcpy(id.bssid, address, AIRPDCAP_MAC_LEN);
-#ifdef _DEBUG
-        sprintf(msgbuf, "BSSID: %2X.%2X.%2X.%2X.%2X.%2X\t", id.bssid[0],id.bssid[1],id.bssid[2],id.bssid[3],id.bssid[4],id.bssid[5]);
-#endif
+    if ( (addr=AirPDcapGetBssidAddress((const AIRPDCAP_MAC_FRAME_ADDR4 *)(data))) != NULL) {
+        memcpy(id.bssid, addr, AIRPDCAP_MAC_LEN);
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", msgbuf, AIRPDCAP_DEBUG_LEVEL_3);
     } else {
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "BSSID not found", AIRPDCAP_DEBUG_LEVEL_5);
         return AIRPDCAP_RET_REQ_DATA;
     }
 
-    /* get STA address */
-    if ( (address=AirPDcapGetStaAddress((const AIRPDCAP_MAC_FRAME_ADDR4 *)(data))) != NULL) {
-        memcpy(id.sta, address, AIRPDCAP_MAC_LEN);
-#ifdef _DEBUG
-        sprintf(msgbuf, "ST_MAC: %2X.%2X.%2X.%2X.%2X.%2X\t", id.sta[0],id.sta[1],id.sta[2],id.sta[3],id.sta[4],id.sta[5]);
-#endif
+    /* get STA address */
+    if ( (addr=AirPDcapGetStaAddress((const AIRPDCAP_MAC_FRAME_ADDR4 *)(data))) != NULL) {
+        memcpy(id.sta, addr, AIRPDCAP_MAC_LEN);
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", msgbuf, AIRPDCAP_DEBUG_LEVEL_3);
     } else {
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "SA not found", AIRPDCAP_DEBUG_LEVEL_5);
         return AIRPDCAP_RET_REQ_DATA;
     }
 
-    /* search for a cached Security Association for current BSSID and station MAC      */
-    if ((index=AirPDcapGetSa(ctx, &id))==-1) {
-        /* create a new Security Association   */
-        if ((index=AirPDcapStoreSa(ctx, &id))==-1) {
-            return AIRPDCAP_RET_UNSUCCESS;
+    /* check if data is encrypted (use the WEP bit in the Frame Control field) */
+    if (AIRPDCAP_WEP(data[1])==0) {
+        if (scanHandshake) {
+            /* data is sent in cleartext, check if is an authentication message or end the process */
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "Unencrypted data", AIRPDCAP_DEBUG_LEVEL_3);
+            return (AirPDcapScanForKeys(ctx, data, mac_header_len, tot_len, id));
         }
-    }
-
-    /* get the Security Association structure  */
-    sa=&ctx->sa[index];
+        return AIRPDCAP_RET_NO_DATA_ENCRYPTED;
+    } else {
+        PAIRPDCAP_SEC_ASSOCIATION sa;
+        int offset = 0;
 
-    /* cache offset in the packet data (to scan encryption data)       */
-    offset = mac_header_len;
+        /* get the Security Association structure for the STA and AP */
+        sa = AirPDcapGetSaPtr(ctx, &id);
+        if (sa == NULL){
+            return AIRPDCAP_RET_REQ_DATA;
+        }
 
-    /* check if data is encrypted (use the WEP bit in the Frame Control field) */
-    if (AIRPDCAP_WEP(data[1])==0)
-    {
-        if (mngHandshake) {
-            /* data is sent in cleartext, check if is an authentication message or end the process     */
-            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "Unencrypted data", AIRPDCAP_DEBUG_LEVEL_3);
+        /* cache offset in the packet data (to scan encryption data) */
+        offset = mac_header_len;
 
-            /* check if the packet as an LLC header and the packet is 802.1X authentication (IEEE 802.1X-2004, pg. 24) */
-            if (memcmp(data+offset, dot1x_header, 8) == 0) {
-                AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "Authentication: EAPOL packet", AIRPDCAP_DEBUG_LEVEL_3);
+        if (decrypt_data==NULL) {
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "no decrypt buffer, use local", AIRPDCAP_DEBUG_LEVEL_3);
+            decrypt_data=tmp_data;
+            decrypt_len=&tmp_len;
+        }
 
-                /* skip LLC header     */
-                offset+=8;
+        /* create new header and data to modify */
+        *decrypt_len = tot_len;
+        memcpy(decrypt_data, data, *decrypt_len);
 
-                /* check the version of the EAPOL protocol used (IEEE 802.1X-2004, pg. 24)     */
-                /* TODO EAPOL protocol version to check?       */
-                /*if (data[offset]!=2) {
-                    AIRPDCAP_DEBUG_PRINT_LINE("EAPOL protocol version not recognized", AIRPDCAP_DEBUG_LEVEL_5);
-                    return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
-                }*/
+        /* encrypted data */
+        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "Encrypted data", AIRPDCAP_DEBUG_LEVEL_3);
 
-                /*     check if the packet is a EAPOL-Key (0x03) (IEEE 802.1X-2004, pg. 25)    */
-                if (data[offset+1]!=3) {
-                    AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "Not EAPOL-Key", AIRPDCAP_DEBUG_LEVEL_5);
-                    return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
-                }
+        /* check the Extension IV to distinguish between WEP encryption and WPA encryption */
+        /* refer to IEEE 802.11i-2004, 8.2.1.2, pag.35 for WEP,    */
+        /*          IEEE 802.11i-2004, 8.3.2.2, pag. 45 for TKIP,  */
+        /*          IEEE 802.11i-2004, 8.3.3.2, pag. 57 for CCMP   */
+        if (AIRPDCAP_EXTIV(data[offset+3])==0) {
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "WEP encryption", AIRPDCAP_DEBUG_LEVEL_3);
+            return AirPDcapWepMng(ctx, decrypt_data, mac_header_len, decrypt_len, key, sa, offset);
+        } else {
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "TKIP or CCMP encryption", AIRPDCAP_DEBUG_LEVEL_3);
 
-                /* get and check the body length (IEEE 802.1X-2004, pg. 25)    */
-                bodyLength=pntohs(data+offset+2);
-                if ((tot_len-offset-4) < bodyLength) {
-                    AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "EAPOL body too short", AIRPDCAP_DEBUG_LEVEL_5);
-                    return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
-                }
+            /* If index >= 1, then use the group key.  This will not work if the AP is using
+               more than one group key simultaneously.  I've not seen this in practice, however.
+               Usually an AP will rotate between the two key index values of 1 and 2 whenever
+               it needs to change the group key to be used. */
+            if (AIRPDCAP_KEY_INDEX(data[offset+3])>=1){
 
-                /* skip EAPOL MPDU and go to the first byte of the body        */
-                offset+=4;
+                AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "The key index >= 1. This is encrypted with a group key.", AIRPDCAP_DEBUG_LEVEL_3);
 
-                /* check if the key descriptor type is valid (IEEE 802.1X-2004, pg. 27)        */
-                if (/*data[offset]!=0x1 &&*/   /* RC4 Key Descriptor Type (deprecated) */
-                    data[offset]!=0x2 &&               /* IEEE 802.11 Key Descriptor Type                      */
-                    data[offset]!=0xFE)                /* TODO what's this value???                                    */
-                {
-                    AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "Not valid key descriptor type", AIRPDCAP_DEBUG_LEVEL_5);
-                    return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
-                }
+                /* force STA address to broadcast MAC so we load the SA for the groupkey */
+                memcpy(id.sta, broadcast_mac, AIRPDCAP_MAC_LEN);
 
-                /* start with descriptor body  */
-                offset+=1;
+#ifdef _DEBUG
+                g_snprintf(msgbuf, MSGBUF_LEN, "ST_MAC: %2X.%2X.%2X.%2X.%2X.%2X\t", id.sta[0],id.sta[1],id.sta[2],id.sta[3],id.sta[4],id.sta[5]);
+                AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", msgbuf, AIRPDCAP_DEBUG_LEVEL_3);
+#endif
 
-                /* manage the 4-way handshake to define the key */
-                return AirPDcapRsna4WHandshake(ctx, data, sa, key, offset);
-            } else {
-                /* cleartext message, not authentication */
-                AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "No authentication data", AIRPDCAP_DEBUG_LEVEL_5);
-                return AIRPDCAP_RET_NO_DATA_ENCRYPTED;
+                /* search for a cached Security Association for current BSSID and broadcast MAC */
+                sa = AirPDcapGetSaPtr(ctx, &id);
+                if (sa == NULL)
+                    return AIRPDCAP_RET_REQ_DATA;
             }
-        }
-    } else {
-        if (mngDecrypt) {
-
-            if (decrypt_data==NULL)
-                return AIRPDCAP_RET_UNSUCCESS;
-
-            /* create new header and data to modify    */
-            *decrypt_len = tot_len;
-            memcpy(decrypt_data, data, *decrypt_len);
-
-            /* encrypted data  */
-            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "Encrypted data", AIRPDCAP_DEBUG_LEVEL_3);
 
-            /* check the Extension IV to distinguish between WEP encryption and WPA encryption */
-            /* refer to IEEE 802.11i-2004, 8.2.1.2, pag.35 for WEP,    */
-            /*         IEEE 802.11i-2004, 8.3.2.2, pag. 45 for TKIP,           */
-            /*         IEEE 802.11i-2004, 8.3.3.2, pag. 57 for CCMP                    */
-            if (AIRPDCAP_EXTIV(data[offset+3])==0) {
-                AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "WEP encryption", AIRPDCAP_DEBUG_LEVEL_3);
-                return AirPDcapWepMng(ctx, decrypt_data, mac_header_len, decrypt_len, key, sa, offset);
-            } else {
-                AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapPacketProcess", "TKIP or CCMP encryption", AIRPDCAP_DEBUG_LEVEL_3);
-                return AirPDcapRsnaMng(decrypt_data, mac_header_len, decrypt_len, key, sa, offset);
+            /* Decrypt the packet using the appropriate SA */
+            if (AirPDcapRsnaMng(decrypt_data, mac_header_len, decrypt_len, key, sa, offset) == AIRPDCAP_RET_SUCCESS) {
+                /* If we successfully decrypted a packet, scan it to see if it contains a key handshake.
+                   The group key handshake could be sent at any time the AP wants to change the key (such as when
+                   it is using key rotation) and it also could be a rekey for the Pairwise key. So we must scan every packet. */
+                if (scanHandshake) {
+                    return (AirPDcapScanForKeys(ctx, decrypt_data, mac_header_len, *decrypt_len, id));
+                } else {
+                    return AIRPDCAP_RET_SUCCESS;
+                }
             }
         }
     }
-
     return AIRPDCAP_RET_UNSUCCESS;
 }
 
@@ -444,17 +790,17 @@ INT AirPDcapSetKeys(
         return 0;
     }
 
-    /* clean keys collection before setting new ones   */
-    AirPDcapCleanKeys(ctx);
+    /* clean key and SA collections before setting new ones */
+    AirPDcapInitContext(ctx);
 
-    /* check and insert keys   */
+    /* check and insert keys */
     for (i=0, success=0; i<(INT)keys_nr; i++) {
         if (AirPDcapValidateKey(keys+i)==TRUE) {
             if (keys[i].KeyType==AIRPDCAP_KEY_TYPE_WPA_PWD) {
                 AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapSetKeys", "Set a WPA-PWD key", AIRPDCAP_DEBUG_LEVEL_4);
                 AirPDcapRsnaPwd2Psk(keys[i].UserPwd.Passphrase, keys[i].UserPwd.Ssid, keys[i].UserPwd.SsidLen, keys[i].KeyData.Wpa.Psk);
             }
-#ifdef _DEBUG
+#ifdef _DEBUG
             else if (keys[i].KeyType==AIRPDCAP_KEY_TYPE_WPA_PMK) {
                 AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapSetKeys", "Set a WPA-PMK key", AIRPDCAP_DEBUG_LEVEL_4);
             } else if (keys[i].KeyType==AIRPDCAP_KEY_TYPE_WEP) {
@@ -474,28 +820,48 @@ INT AirPDcapSetKeys(
     return success;
 }
 
-INT AirPDcapCleanKeys(
+static void
+AirPDcapCleanKeys(
     PAIRPDCAP_CONTEXT ctx)
 {
-    INT i;
     AIRPDCAP_DEBUG_TRACE_START("AirPDcapCleanKeys");
 
     if (ctx==NULL) {
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapCleanKeys", "NULL context", AIRPDCAP_DEBUG_LEVEL_5);
         AIRPDCAP_DEBUG_TRACE_END("AirPDcapCleanKeys");
-        return 0;
+        return;
     }
 
-    for (i=0; i<AIRPDCAP_MAX_KEYS_NR; i++) {
-        memset(&ctx->keys[i], 0, sizeof(AIRPDCAP_KEY_ITEM));
-    }
+    memset(ctx->keys, 0, sizeof(AIRPDCAP_KEY_ITEM) * AIRPDCAP_MAX_KEYS_NR);
 
     ctx->keys_nr=0;
 
     AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapCleanKeys", "Keys collection cleaned!", AIRPDCAP_DEBUG_LEVEL_5);
     AIRPDCAP_DEBUG_TRACE_END("AirPDcapCleanKeys");
+}
+
+static void
+AirPDcapRecurseCleanSA(
+    PAIRPDCAP_SEC_ASSOCIATION sa)
+{
+    if (sa->next != NULL) {
+        AirPDcapRecurseCleanSA(sa->next);
+        g_free(sa->next);
+        sa->next = NULL;
+    }
+}
+
+static void
+AirPDcapCleanSecAssoc(
+    PAIRPDCAP_CONTEXT ctx)
+{
+    PAIRPDCAP_SEC_ASSOCIATION psa;
+    int i;
 
-    return i;
+    for (psa = ctx->sa, i = 0; i < AIRPDCAP_MAX_SEC_ASSOCIATIONS_NR; i++, psa++) {
+        /* To iterate is human, to recurse, divine */
+        AirPDcapRecurseCleanSA(psa);
+    }
 }
 
 INT AirPDcapGetKeys(
@@ -562,9 +928,11 @@ INT AirPDcapInitContext(
 
     ctx->first_free_index=0;
     ctx->index=-1;
-    ctx->last_stored_index=-1;
+    ctx->sa_index=-1;
     ctx->pkt_ssid_len = 0;
 
+    memset(ctx->sa, 0, AIRPDCAP_MAX_SEC_ASSOCIATIONS_NR * sizeof(AIRPDCAP_SEC_ASSOCIATION));
+
     AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapInitContext", "Context initialized!", AIRPDCAP_DEBUG_LEVEL_5);
     AIRPDCAP_DEBUG_TRACE_END("AirPDcapInitContext");
     return AIRPDCAP_RET_SUCCESS;
@@ -582,26 +950,27 @@ INT AirPDcapDestroyContext(
     }
 
     AirPDcapCleanKeys(ctx);
+    AirPDcapCleanSecAssoc(ctx);
 
     ctx->first_free_index=0;
     ctx->index=-1;
-    ctx->last_stored_index=-1;
+    ctx->sa_index=-1;
 
     AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapDestroyContext", "Context destroyed!", AIRPDCAP_DEBUG_LEVEL_5);
     AIRPDCAP_DEBUG_TRACE_END("AirPDcapDestroyContext");
     return AIRPDCAP_RET_SUCCESS;
 }
 
-#ifdef __cplusplus
+#ifdef __cplusplus
 }
 #endif
 
 /****************************************************************************/
 
 /****************************************************************************/
-/* Internal function definitions                                               */
+/* Internal function definitions                                         */
 
-#ifdef __cplusplus
+#ifdef __cplusplus
 extern "C" {
 #endif
 
@@ -614,52 +983,94 @@ AirPDcapRsnaMng(
     AIRPDCAP_SEC_ASSOCIATION *sa,
     INT offset)
 {
-    INT ret_value;
+    INT ret_value=1;
+    UCHAR *try_data;
+    guint try_data_len = *decrypt_len;
 
     if (sa->key==NULL) {
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "No key associated", AIRPDCAP_DEBUG_LEVEL_3);
         return AIRPDCAP_RET_REQ_DATA;
     }
-    if (sa->validKey==FALSE) {
-        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "Key not yet valid", AIRPDCAP_DEBUG_LEVEL_3);
+
+    if (*decrypt_len > try_data_len) {
+        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "Invalid decryption length", AIRPDCAP_DEBUG_LEVEL_3);
         return AIRPDCAP_RET_UNSUCCESS;
     }
-    if (sa->wpa.key_ver==1) {
-        /*     CCMP -> HMAC-MD5 is the EAPOL-Key MIC, RC4 is the EAPOL-Key encryption algorithm        */
-        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "TKIP", AIRPDCAP_DEBUG_LEVEL_3);
-
-        ret_value=AirPDcapTkipDecrypt(decrypt_data+offset, *decrypt_len-offset, decrypt_data+AIRPDCAP_TA_OFFSET, AIRPDCAP_GET_TK(sa->wpa.ptk));
-        if (ret_value)
-            return ret_value;
 
-        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "TKIP DECRYPTED!!!", AIRPDCAP_DEBUG_LEVEL_3);
-        /* remove MIC (8bytes) and ICV (4bytes) from the end of packet */
-        *decrypt_len-=12;
-    } else {
-        /*     AES-CCMP -> HMAC-SHA1-128 is the EAPOL-Key MIC, AES wep_key wrap is the EAPOL-Key encryption algorithm  */
-        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "CCMP", AIRPDCAP_DEBUG_LEVEL_3);
+    /* allocate a temp buffer for the decryption loop */
+    try_data=(UCHAR *)g_malloc(try_data_len);
+
+    /* start of loop added by GCS */
+    for(/* sa */; sa != NULL ;sa=sa->next) {
+
+       if (sa->validKey==FALSE) {
+           AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "Key not yet valid", AIRPDCAP_DEBUG_LEVEL_3);
+           continue;
+       }
+
+       /* copy the encrypted data into a temp buffer */
+       memcpy(try_data, decrypt_data, *decrypt_len);
+
+       if (sa->wpa.key_ver==1) {
+           /* CCMP -> HMAC-MD5 is the EAPOL-Key MIC, RC4 is the EAPOL-Key encryption algorithm */
+           AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "TKIP", AIRPDCAP_DEBUG_LEVEL_3);
+           DEBUG_DUMP("ptk", sa->wpa.ptk, 64);
+           DEBUG_DUMP("ptk portion used", AIRPDCAP_GET_TK(sa->wpa.ptk), 16);
+
+           ret_value=AirPDcapTkipDecrypt(try_data+offset, *decrypt_len-offset, try_data+AIRPDCAP_TA_OFFSET, AIRPDCAP_GET_TK(sa->wpa.ptk));
+           if (ret_value){
+               AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "TKIP failed!", AIRPDCAP_DEBUG_LEVEL_3);
+               continue;
+           }
+
+           AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "TKIP DECRYPTED!!!", AIRPDCAP_DEBUG_LEVEL_3);
+           /* remove MIC (8bytes) and ICV (4bytes) from the end of packet */
+           *decrypt_len-=12;
+           break;
+       } else {
+           /* AES-CCMP -> HMAC-SHA1-128 is the EAPOL-Key MIC, AES wep_key wrap is the EAPOL-Key encryption algorithm */
+           AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "CCMP", AIRPDCAP_DEBUG_LEVEL_3);
+
+           ret_value=AirPDcapCcmpDecrypt(try_data, mac_header_len, (INT)*decrypt_len, AIRPDCAP_GET_TK(sa->wpa.ptk));
+           if (ret_value)
+              continue;
+
+           AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "CCMP DECRYPTED!!!", AIRPDCAP_DEBUG_LEVEL_3);
+           /* remove MIC (8bytes) from the end of packet */
+           *decrypt_len-=8;
+           break;
+       }
+    }
+    /* end of loop */
 
-        ret_value=AirPDcapCcmpDecrypt(decrypt_data, mac_header_len, (INT)*decrypt_len, AIRPDCAP_GET_TK(sa->wpa.ptk));
-        if (ret_value)
-            return ret_value;
+    /* none of the keys worked */
+    if(sa == NULL) {
+        g_free(try_data);
+        return ret_value;
+    }
 
-        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "CCMP DECRYPTED!!!", AIRPDCAP_DEBUG_LEVEL_3);
-        /* remove MIC (8bytes) from the end of packet  */
-        *decrypt_len-=8;
+    if (*decrypt_len > try_data_len || *decrypt_len < 8) {
+        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsnaMng", "Invalid decryption length", AIRPDCAP_DEBUG_LEVEL_3);
+        g_free(try_data);
+        return AIRPDCAP_RET_UNSUCCESS;
     }
 
-    /* remove protection bit   */
+    /* copy the decrypted data into the decrypt buffer GCS*/
+    memcpy(decrypt_data, try_data, *decrypt_len);
+    g_free(try_data);
+
+    /* remove protection bit */
     decrypt_data[1]&=0xBF;
 
-    /* remove TKIP/CCMP header */
+    /* remove TKIP/CCMP header */
     offset = mac_header_len;
     *decrypt_len-=8;
-    memcpy(decrypt_data+offset, decrypt_data+offset+8, *decrypt_len-offset);
+    memmove(decrypt_data+offset, decrypt_data+offset+8, *decrypt_len-offset);
 
     if (key!=NULL) {
         memcpy(key, sa->key, sizeof(AIRPDCAP_KEY_ITEM));
-
-        if (sa->wpa.key_ver==AIRPDCAP_WPA_KEY_VER_CCMP)
+        memcpy(key->KeyData.Wpa.Ptk, sa->wpa.ptk, AIRPDCAP_WPA_PTK_LEN); /* copy the PTK to the key structure for future use by wireshark */
+        if (sa->wpa.key_ver==AIRPDCAP_WPA_KEY_VER_NOT_CCMP)
             key->KeyType=AIRPDCAP_KEY_TYPE_TKIP;
         else if (sa->wpa.key_ver==AIRPDCAP_WPA_KEY_VER_AES_CCMP)
             key->KeyType=AIRPDCAP_KEY_TYPE_CCMP;
@@ -684,13 +1095,16 @@ AirPDcapWepMng(
     INT key_index;
     AIRPDCAP_KEY_ITEM *tmp_key;
     UINT8 useCache=FALSE;
-    UCHAR *try_data = ep_alloc(*decrypt_len);
+    UCHAR *try_data;
+    guint try_data_len = *decrypt_len;
+
+    try_data = (UCHAR *)g_malloc(try_data_len);
 
     if (sa->key!=NULL)
         useCache=TRUE;
 
     for (key_index=0; key_index<(INT)ctx->keys_nr; key_index++) {
-        /* use the cached one, or try all keys */
+        /* use the cached one, or try all keys */
         if (!useCache) {
             tmp_key=&ctx->keys[key_index];
         } else {
@@ -703,15 +1117,14 @@ AirPDcapWepMng(
             }
         }
 
-        /* obviously, try only WEP keys...     */
-        if (tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WEP)
-        {
+        /* obviously, try only WEP keys... */
+        if (tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WEP) {
             AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapWepMng", "Try WEP key...", AIRPDCAP_DEBUG_LEVEL_3);
 
             memset(wep_key, 0, sizeof(wep_key));
             memcpy(try_data, decrypt_data, *decrypt_len);
 
-            /* Costruct the WEP seed: copy the IV in first 3 bytes and then the WEP key (refer to 802-11i-2004, 8.2.1.4.3, pag. 36)    */
+            /* Costruct the WEP seed: copy the IV in first 3 bytes and then the WEP key (refer to 802-11i-2004, 8.2.1.4.3, pag. 36) */
             memcpy(wep_key, try_data+mac_header_len, AIRPDCAP_WEP_IVLEN);
             keylen=tmp_key->KeyData.Wep.WepKeyLen;
             memcpy(wep_key+AIRPDCAP_WEP_IVLEN, tmp_key->KeyData.Wep.WepKey, keylen);
@@ -726,18 +1139,18 @@ AirPDcapWepMng(
         }
 
         if (!ret_value && tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WEP) {
-            /* the tried key is the correct one, cached in the Security Association    */
+            /* the tried key is the correct one, cached in the Security Association */
 
             sa->key=tmp_key;
 
             if (key!=NULL) {
-                memcpy(key, &sa->key, sizeof(AIRPDCAP_KEY_ITEM));
+                memcpy(key, sa->key, sizeof(AIRPDCAP_KEY_ITEM));
                 key->KeyType=AIRPDCAP_KEY_TYPE_WEP;
             }
 
             break;
         } else {
-            /* the cached key was not valid, try other keys    */
+            /* the cached key was not valid, try other keys */
 
             if (useCache==TRUE) {
                 useCache=FALSE;
@@ -746,35 +1159,41 @@ AirPDcapWepMng(
         }
     }
 
+    g_free(try_data);
     if (ret_value)
-        return ret_value;
+        return AIRPDCAP_RET_UNSUCCESS;
 
     AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapWepMng", "WEP DECRYPTED!!!", AIRPDCAP_DEBUG_LEVEL_3);
 
-    /* remove ICV (4bytes) from the end of packet      */
+    /* remove ICV (4bytes) from the end of packet */
     *decrypt_len-=4;
 
-    /* remove protection bit   */
+    if (*decrypt_len < 4) {
+        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapWepMng", "Decryption length too short", AIRPDCAP_DEBUG_LEVEL_3);
+        return AIRPDCAP_RET_UNSUCCESS;
+    }
+
+    /* remove protection bit */
     decrypt_data[1]&=0xBF;
 
-    /* remove IC header        */
+    /* remove IC header */
     offset = mac_header_len;
     *decrypt_len-=4;
-    memcpy(decrypt_data+offset, decrypt_data+offset+AIRPDCAP_WEP_IVLEN+AIRPDCAP_WEP_KIDLEN, *decrypt_len-offset);
+    memmove(decrypt_data+offset, decrypt_data+offset+AIRPDCAP_WEP_IVLEN+AIRPDCAP_WEP_KIDLEN, *decrypt_len-offset);
 
     return AIRPDCAP_RET_SUCCESS;
 }
 
-/* Refer to IEEE 802.11i-2004, 8.5.3, pag. 85  */
+/* Refer to IEEE 802.11i-2004, 8.5.3, pag. 85 */
 static INT
 AirPDcapRsna4WHandshake(
     PAIRPDCAP_CONTEXT ctx,
     const UCHAR *data,
     AIRPDCAP_SEC_ASSOCIATION *sa,
-    PAIRPDCAP_KEY_ITEM key,
     INT offset)
 {
-    AIRPDCAP_KEY_ITEM *tmp_key, pkt_key;
+    AIRPDCAP_KEY_ITEM *tmp_key, *tmp_pkt_key, pkt_key;
+    AIRPDCAP_SEC_ASSOCIATION *tmp_sa;
     INT key_index;
     INT ret_value=1;
     UCHAR useCache=FALSE;
@@ -784,38 +1203,46 @@ AirPDcapRsna4WHandshake(
     if (sa->key!=NULL)
         useCache=TRUE;
 
-    /* a 4-way handshake packet use a Pairwise key type (IEEE 802.11i-2004, pg. 79)    */
+    /* a 4-way handshake packet use a Pairwise key type (IEEE 802.11i-2004, pg. 79) */
     if (AIRPDCAP_EAP_KEY(data[offset+1])!=1) {
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "Group/STAKey message (not used)", AIRPDCAP_DEBUG_LEVEL_5);
         return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
     }
 
-    /* TODO timeouts? reauthentication?        */
+    /* TODO timeouts? */
 
-    /* TODO consider key-index */
+    /* TODO consider key-index */
 
-    /* TODO considera Deauthentications        */
+    /* TODO considera Deauthentications */
 
     AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "4-way handshake...", AIRPDCAP_DEBUG_LEVEL_5);
 
-    /* manage 4-way handshake packets; this step completes the 802.1X authentication process (IEEE 802.11i-2004, pag. 85)      */
+    /* manage 4-way handshake packets; this step completes the 802.1X authentication process (IEEE 802.11i-2004, pag. 85) */
 
-    /* message 1: Authenticator->Supplicant (Sec=0, Mic=0, Ack=1, Inst=0, Key=1(pairwise), KeyRSC=0, Nonce=ANonce, MIC=0)      */
+    /* message 1: Authenticator->Supplicant (Sec=0, Mic=0, Ack=1, Inst=0, Key=1(pairwise), KeyRSC=0, Nonce=ANonce, MIC=0) */
     if (AIRPDCAP_EAP_INST(data[offset+1])==0 &&
         AIRPDCAP_EAP_ACK(data[offset+1])==1 &&
         AIRPDCAP_EAP_MIC(data[offset])==0)
     {
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "4-way handshake message 1", AIRPDCAP_DEBUG_LEVEL_3);
 
-        /* On reception of Message 1, the Supplicant determines whether the Key Replay Counter field value has been                    */
-        /* used before with the current PMKSA. If the Key Replay Counter field value is less than or equal to the current      */
-        /* local value, the Supplicant discards the message.                                                                                                                                                                   */
-        /* -> not checked, the Authenticator will be send another Message 1 (hopefully!)                                                                                               */
+        /* On reception of Message 1, the Supplicant determines whether the Key Replay Counter field value has been        */
+        /* used before with the current PMKSA. If the Key Replay Counter field value is less than or equal to the current  */
+        /* local value, the Supplicant discards the message.                                                               */
+        /* -> not checked, the Authenticator will be send another Message 1 (hopefully!)                                   */
+
+        /* This saves the sa since we are reauthenticating which will overwrite our current sa GCS*/
+        if( sa->handshake >= 2) {
+            tmp_sa= g_new(AIRPDCAP_SEC_ASSOCIATION, 1);
+            memcpy(tmp_sa, sa, sizeof(AIRPDCAP_SEC_ASSOCIATION));
+            sa->validKey=FALSE;
+            sa->next=tmp_sa;
+        }
 
-        /* save ANonce (from authenticator)    to derive the PTK with the SNonce (from the 2 message)  */
+        /* save ANonce (from authenticator) to derive the PTK with the SNonce (from the 2 message) */
         memcpy(sa->wpa.nonce, data+offset+12, 32);
 
-        /* get the Key Descriptor Version (to select algorithm used in decryption -CCMP or TKIP-)      */
+        /* get the Key Descriptor Version (to select algorithm used in decryption -CCMP or TKIP-) */
         sa->wpa.key_ver=AIRPDCAP_EAP_KEY_DESCR_VER(data[offset+1]);
 
         sa->handshake=1;
@@ -823,177 +1250,163 @@ AirPDcapRsna4WHandshake(
         return AIRPDCAP_RET_SUCCESS_HANDSHAKE;
     }
 
-    /* message 2|4: Supplicant->Authenticator (Sec=0|1, Mic=1, Ack=0, Inst=0, Key=1(pairwise), KeyRSC=0, Nonce=SNonce|0, MIC=MIC(KCK,EAPOL))   */
+    /* message 2|4: Supplicant->Authenticator (Sec=0|1, Mic=1, Ack=0, Inst=0, Key=1(pairwise), KeyRSC=0, Nonce=SNonce|0, MIC=MIC(KCK,EAPOL)) */
     if (AIRPDCAP_EAP_INST(data[offset+1])==0 &&
         AIRPDCAP_EAP_ACK(data[offset+1])==0 &&
         AIRPDCAP_EAP_MIC(data[offset])==1)
     {
-        if (AIRPDCAP_EAP_SEC(data[offset])==0) {
-
-            /* PATCH:  some implementations set secure bit to 0 also in the 4th message                */
-            /*         to recognize which message is this check if wep_key data length is 0            */
-            /*         in the 4th message                                                              */
-            if (data[offset+92]!=0 || data[offset+93]!=0) {
-                /* message 2   */
-                AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "4-way handshake message 2", AIRPDCAP_DEBUG_LEVEL_3);
-
-                /* On reception of Message 2, the Authenticator checks that the key replay counter corresponds to the  */
-                /* outstanding Message 1. If not, it silently discards the message.                                                                                            */
-                /* If the calculated MIC does not match the MIC that the Supplicant included in the EAPOL-Key frame,   */
-                /* the Authenticator silently discards Message 2.                                                                                                                                              */
-                /* -> not checked; the Supplicant will send another message 2 (hopefully!)                                                                             */
-
-                /* now you can derive the PTK  */
-                for (key_index=0; key_index<(INT)ctx->keys_nr || useCache; key_index++) {
-                    /* use the cached one, or try all keys     */
-                    if (!useCache) {
-                        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "Try WPA key...", AIRPDCAP_DEBUG_LEVEL_3);
-                        tmp_key=&ctx->keys[key_index];
+         /* Check nonce to differentiate between message 2 or 4
+          * nonce will be non zero for message 2 and zero for message 4.
+          * At least needed for Windows, since it is setting the secure bit on message 2 when rekeying */
+        if (!memiszero(data+offset+12, 32)) {
+            /* message 2 */
+            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "4-way handshake message 2", AIRPDCAP_DEBUG_LEVEL_3);
+
+            /* On reception of Message 2, the Authenticator checks that the key replay counter corresponds to the */
+            /* outstanding Message 1. If not, it silently discards the message.                                   */
+            /* If the calculated MIC does not match the MIC that the Supplicant included in the EAPOL-Key frame,  */
+            /* the Authenticator silently discards Message 2.                                                     */
+            /* -> not checked; the Supplicant will send another message 2 (hopefully!)                            */
+
+            /* now you can derive the PTK */
+            for (key_index=0; key_index<(INT)ctx->keys_nr || useCache; key_index++) {
+                /* use the cached one, or try all keys */
+                if (!useCache) {
+                    AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "Try WPA key...", AIRPDCAP_DEBUG_LEVEL_3);
+                    tmp_key=&ctx->keys[key_index];
+                } else {
+                    /* there is a cached key in the security association, if it's a WPA key try it... */
+                    if (sa->key!=NULL &&
+                        (sa->key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PWD ||
+                         sa->key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PSK ||
+                         sa->key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PMK)) {
+                            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "Try cached WPA key...", AIRPDCAP_DEBUG_LEVEL_3);
+                            tmp_key=sa->key;
                     } else {
-                        /* there is a cached key in the security association, if it's a WPA key try it...      */
-                        if (sa->key!=NULL &&
-                            (sa->key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PWD ||
-                             sa->key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PSK ||
-                             sa->key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PMK)) {
-                                AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "Try cached WPA key...", AIRPDCAP_DEBUG_LEVEL_3);
-                                tmp_key=sa->key;
-                        } else {
-                            AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "Cached key is of a wrong type, try WPA key...", AIRPDCAP_DEBUG_LEVEL_3);
-                            tmp_key=&ctx->keys[key_index];
-                        }
-                    }
-
-                    /* obviously, try only WPA keys... */
-                    if (tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PWD ||
-                        tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PSK ||
-                        tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PMK)
-                    {
-                        if (tmp_key->KeyType == AIRPDCAP_KEY_TYPE_WPA_PWD && tmp_key->UserPwd.SsidLen == 0 && ctx->pkt_ssid_len > 0 && ctx->pkt_ssid_len <= AIRPDCAP_WPA_SSID_MAX_LEN) {
-                            /* We have a "wildcard" SSID.  Use the one from the packet. */
-                            memcpy(&pkt_key, tmp_key, sizeof(pkt_key));
-                            memcpy(&pkt_key.UserPwd.Ssid, ctx->pkt_ssid, ctx->pkt_ssid_len);
-                             pkt_key.UserPwd.SsidLen = ctx->pkt_ssid_len;
-                            AirPDcapRsnaPwd2Psk(pkt_key.UserPwd.Passphrase, pkt_key.UserPwd.Ssid,
-                                pkt_key.UserPwd.SsidLen, pkt_key.KeyData.Wpa.Psk);
-                            tmp_key = &pkt_key;
-                        }
-
-                        /* derive the PTK from the BSSID, STA MAC, PMK, SNonce, ANonce */
-                        AirPDcapRsnaPrfX(sa,                   /* authenticator nonce, bssid, station mac      */
-                            tmp_key->KeyData.Wpa.Pmk,  /* PMK  */
-                            data+offset+12,            /* supplicant nonce     */
-                            512,
-                            sa->wpa.ptk);
-
-                        /* verify the MIC (compare the MIC in the packet included in this message with a MIC calculated with the PTK)  */
-                        eapol_len=pntohs(data+offset-3)+4;
-                        memcpy(eapol, &data[offset-5], (eapol_len<AIRPDCAP_EAPOL_MAX_LEN?eapol_len:AIRPDCAP_EAPOL_MAX_LEN));
-                        ret_value=AirPDcapRsnaMicCheck(eapol,                                          /*      eapol frame (header also)               */
-                            eapol_len,                                                                                                 /*      eapol frame length                              */
-                            sa->wpa.ptk,                                                                                               /*      Key Confirmation Key                            */
-                            AIRPDCAP_EAP_KEY_DESCR_VER(data[offset+1]));               /*      EAPOL-Key description version   */
-
-                        /* If the MIC is valid, the Authenticator checks that the RSN information element bit-wise matches             */
-                        /* that from the (Re)Association Request message.                                                                                                                                              */
-                        /*             i) TODO If these are not exactly the same, the Authenticator uses MLME-DEAUTHENTICATE.request   */
-                        /* primitive to terminate the association.                                                                                                                                                             */
-                        /*             ii) If they do match bit-wise, the Authenticator constructs Message 3.                                                                  */
+                        AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "Cached key is of a wrong type, try WPA key...", AIRPDCAP_DEBUG_LEVEL_3);
+                        tmp_key=&ctx->keys[key_index];
                     }
+                }
 
-                    if (!ret_value &&
-                        (tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PWD ||
-                        tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PSK ||
-                        tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PMK))
-                    {
-                        /* the temporary key is the correct one, cached in the Security Association    */
-
-                        sa->key=tmp_key;
-
-                        if (key!=NULL) {
-                            memcpy(key, &tmp_key, sizeof(AIRPDCAP_KEY_ITEM));
-                            if (AIRPDCAP_EAP_KEY_DESCR_VER(data[offset+1])==AIRPDCAP_WPA_KEY_VER_CCMP)
-                                key->KeyType=AIRPDCAP_KEY_TYPE_TKIP;
-                            else if (AIRPDCAP_EAP_KEY_DESCR_VER(data[offset+1])==AIRPDCAP_WPA_KEY_VER_AES_CCMP)
-                                key->KeyType=AIRPDCAP_KEY_TYPE_CCMP;
-                        }
-
-                        break;
+                /* obviously, try only WPA keys... */
+                if (tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PWD ||
+                    tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PSK ||
+                    tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PMK)
+                {
+                    if (tmp_key->KeyType == AIRPDCAP_KEY_TYPE_WPA_PWD && tmp_key->UserPwd.SsidLen == 0 && ctx->pkt_ssid_len > 0 && ctx->pkt_ssid_len <= AIRPDCAP_WPA_SSID_MAX_LEN) {
+                        /* We have a "wildcard" SSID.  Use the one from the packet. */
+                        memcpy(&pkt_key, tmp_key, sizeof(pkt_key));
+                        memcpy(&pkt_key.UserPwd.Ssid, ctx->pkt_ssid, ctx->pkt_ssid_len);
+                         pkt_key.UserPwd.SsidLen = ctx->pkt_ssid_len;
+                        AirPDcapRsnaPwd2Psk(pkt_key.UserPwd.Passphrase, pkt_key.UserPwd.Ssid,
+                            pkt_key.UserPwd.SsidLen, pkt_key.KeyData.Wpa.Psk);
+                        tmp_pkt_key = &pkt_key;
                     } else {
-                        /* the cached key was not valid, try other keys        */
-
-                        if (useCache==TRUE) {
-                            useCache=FALSE;
-                            key_index--;
-                        }
+                        tmp_pkt_key = tmp_key;
                     }
-                }
 
-                if (ret_value) {
-                    AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "handshake step failed", AIRPDCAP_DEBUG_LEVEL_3);
-                    return ret_value;
+                    /* derive the PTK from the BSSID, STA MAC, PMK, SNonce, ANonce */
+                    AirPDcapRsnaPrfX(sa,                            /* authenticator nonce, bssid, station mac */
+                                     tmp_pkt_key->KeyData.Wpa.Psk,      /* PSK == PMK */
+                                     data+offset+12,                /* supplicant nonce */
+                                     512,
+                                     sa->wpa.ptk);
+
+                    /* verify the MIC (compare the MIC in the packet included in this message with a MIC calculated with the PTK) */
+                    eapol_len=pntoh16(data+offset-3)+4;
+                    memcpy(eapol, &data[offset-5], (eapol_len<AIRPDCAP_EAPOL_MAX_LEN?eapol_len:AIRPDCAP_EAPOL_MAX_LEN));
+                    ret_value=AirPDcapRsnaMicCheck(eapol,           /*      eapol frame (header also) */
+                                                   eapol_len,       /*      eapol frame length        */
+                                                   sa->wpa.ptk,     /*      Key Confirmation Key      */
+                                                   AIRPDCAP_EAP_KEY_DESCR_VER(data[offset+1])); /*  EAPOL-Key description version */
+
+                    /* If the MIC is valid, the Authenticator checks that the RSN information element bit-wise matches       */
+                    /* that from the (Re)Association Request message.                                                        */
+                    /*              i) TODO If these are not exactly the same, the Authenticator uses MLME-DEAUTHENTICATE.request */
+                    /* primitive to terminate the association.                                                               */
+                    /*              ii) If they do match bit-wise, the Authenticator constructs Message 3.                   */
                 }
 
-                sa->handshake=2;
-
-                return AIRPDCAP_RET_SUCCESS_HANDSHAKE;
-            } else {
-                /* message 4   */
-
-                /* TODO "Note that when the 4-Way Handshake is first used Message 4 is sent in the clear."     */
+                if (!ret_value &&
+                    (tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PWD ||
+                    tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PSK ||
+                    tmp_key->KeyType==AIRPDCAP_KEY_TYPE_WPA_PMK))
+                {
+                    /* the temporary key is the correct one, cached in the Security Association */
 
-                /* TODO check MIC and Replay Counter                                                                                                                                                                                   */
-                /* On reception of Message 4, the Authenticator verifies that the Key Replay Counter field value is one        */
-                /* that it used on this 4-Way Handshake; if it is not, it silently discards the message.                                               */
-                /* If the calculated MIC does not match the MIC that the Supplicant included in the EAPOL-Key frame, the       */
-                /* Authenticator silently discards Message 4.                                                                                                                                                          */
+                    sa->key=tmp_key;
+                    break;
+                } else {
+                    /* the cached key was not valid, try other keys */
 
-                AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "4-way handshake message 4 (patched)", AIRPDCAP_DEBUG_LEVEL_3);
+                    if (useCache==TRUE) {
+                        useCache=FALSE;
+                        key_index--;
+                    }
+                }
+            }
 
-                sa->handshake=4;
+            if (ret_value) {
+                AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "handshake step failed", AIRPDCAP_DEBUG_LEVEL_3);
+                return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
+            }
 
-                sa->validKey=TRUE;
+            sa->handshake=2;
+            sa->validKey=TRUE; /* we can use the key to decode, even if we have not captured the other eapol packets */
 
-                return AIRPDCAP_RET_SUCCESS_HANDSHAKE;
-            }
-            /* END OF PATCH                                                                                                                                                                    */
-            /*                                                                                                                                                                                                         */
+            return AIRPDCAP_RET_SUCCESS_HANDSHAKE;
         } else {
-            /* message 4       */
+        /* message 4 */
 
-            /* TODO "Note that when the 4-Way Handshake is first used Message 4 is sent in the clear." */
+            /* TODO "Note that when the 4-Way Handshake is first used Message 4 is sent in the clear." */
 
-            /* TODO check MIC and Replay Counter                                                                                                                                                                                       */
-            /* On reception of Message 4, the Authenticator verifies that the Key Replay Counter field value is one    */
-            /* that it used on this 4-Way Handshake; if it is not, it silently discards the message.                                           */
-            /* If the calculated MIC does not match the MIC that the Supplicant included in the EAPOL-Key frame, the   */
-            /* Authenticator silently discards Message 4.                                                                                                                                                              */
+            /* TODO check MIC and Replay Counter                                                                     */
+            /* On reception of Message 4, the Authenticator verifies that the Key Replay Counter field value is one  */
+            /* that it used on this 4-Way Handshake; if it is not, it silently discards the message.                 */
+            /* If the calculated MIC does not match the MIC that the Supplicant included in the EAPOL-Key frame, the */
+            /* Authenticator silently discards Message 4.                                                            */
 
             AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "4-way handshake message 4", AIRPDCAP_DEBUG_LEVEL_3);
 
             sa->handshake=4;
 
-            sa->validKey=TRUE;
-
             return AIRPDCAP_RET_SUCCESS_HANDSHAKE;
         }
     }
 
-    /* message 3: Authenticator->Supplicant (Sec=1, Mic=1, Ack=1, Inst=0/1, Key=1(pairwise), KeyRSC=???, Nonce=ANonce, MIC=1)  */
+    /* message 3: Authenticator->Supplicant (Sec=1, Mic=1, Ack=1, Inst=0/1, Key=1(pairwise), KeyRSC=???, Nonce=ANonce, MIC=1) */
     if (AIRPDCAP_EAP_ACK(data[offset+1])==1 &&
         AIRPDCAP_EAP_MIC(data[offset])==1)
     {
+        const EAPOL_RSN_KEY *pEAPKey;
         AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapRsna4WHandshake", "4-way handshake message 3", AIRPDCAP_DEBUG_LEVEL_3);
 
-        /* On reception of Message 3, the Supplicant silently discards the message if the Key Replay Counter field             */
-        /* value has already been used or if the ANonce value in Message 3 differs from the ANonce value in Message 1. */
-        /* -> not checked, the Authenticator will send another message 3 (hopefully!)                                                                                          */
+        /* On reception of Message 3, the Supplicant silently discards the message if the Key Replay Counter field     */
+        /* value has already been used or if the ANonce value in Message 3 differs from the ANonce value in Message 1. */
+        /* -> not checked, the Authenticator will send another message 3 (hopefully!)                                  */
 
-        /*     TODO check page 88 (RNS)        */
+        /* TODO check page 88 (RNS) */
 
-        return AIRPDCAP_RET_SUCCESS_HANDSHAKE;
+        /* If using WPA2 PSK, message 3 will contain an RSN for the group key (GTK KDE).
+           In order to properly support decrypting WPA2-PSK packets, we need to parse this to get the group key. */
+        pEAPKey = (const EAPOL_RSN_KEY *)(&(data[offset-1]));
+        if (pEAPKey->type == AIRPDCAP_RSN_WPA2_KEY_DESCRIPTOR){
+            PAIRPDCAP_SEC_ASSOCIATION broadcast_sa;
+            AIRPDCAP_SEC_ASSOCIATION_ID id;
+
+            /* Get broadcacst SA for the current BSSID */
+            memcpy(id.sta, broadcast_mac, AIRPDCAP_MAC_LEN);
+            memcpy(id.bssid, sa->saId.bssid, AIRPDCAP_MAC_LEN);
+            broadcast_sa = AirPDcapGetSaPtr(ctx, &id);
+
+            if (broadcast_sa == NULL){
+                return AIRPDCAP_RET_REQ_DATA;
+            }
+            return (AirPDcapDecryptWPABroadcastKey(pEAPKey, sa->wpa.ptk+16, broadcast_sa, FALSE));
+        }
     }
 
-    return AIRPDCAP_RET_UNSUCCESS;
+    return AIRPDCAP_RET_NO_VALID_HANDSHAKE;
 }
 
 static INT
@@ -1004,16 +1417,16 @@ AirPDcapRsnaMicCheck(
     USHORT key_ver)
 {
     UCHAR mic[AIRPDCAP_WPA_MICKEY_LEN];
-    UCHAR c_mic[20];   /* MIC 16 byte, the HMAC-SHA1 use a buffer of 20 bytes */
+    UCHAR c_mic[20];  /* MIC 16 byte, the HMAC-SHA1 use a buffer of 20 bytes */
 
-    /* copy the MIC from the EAPOL packet      */
+    /* copy the MIC from the EAPOL packet */
     memcpy(mic, eapol+AIRPDCAP_WPA_MICKEY_OFFSET+4, AIRPDCAP_WPA_MICKEY_LEN);
 
     /* set to 0 the MIC in the EAPOL packet (to calculate the MIC) */
     memset(eapol+AIRPDCAP_WPA_MICKEY_OFFSET+4, 0, AIRPDCAP_WPA_MICKEY_LEN);
 
-    if (key_ver==AIRPDCAP_WPA_KEY_VER_CCMP) {
-        /* use HMAC-MD5 for the EAPOL-Key MIC  */
+    if (key_ver==AIRPDCAP_WPA_KEY_VER_NOT_CCMP) {
+        /* use HMAC-MD5 for the EAPOL-Key MIC */
         md5_hmac(eapol, eapol_len, KCK, AIRPDCAP_WPA_KCK_LEN, c_mic);
     } else if (key_ver==AIRPDCAP_WPA_KEY_VER_AES_CCMP) {
         /* use HMAC-SHA1-128 for the EAPOL-Key MIC */
@@ -1042,7 +1455,7 @@ AirPDcapValidateKey(
 
     switch (key->KeyType) {
         case AIRPDCAP_KEY_TYPE_WEP:
-            /* check key size limits   */
+            /* check key size limits */
             len=key->KeyData.Wep.WepKeyLen;
             if (len<AIRPDCAP_WEP_KEY_MINLEN || len>AIRPDCAP_WEP_KEY_MAXLEN) {
                 AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapValidateKey", "WEP key: key length not accepted", AIRPDCAP_DEBUG_LEVEL_5);
@@ -1051,19 +1464,19 @@ AirPDcapValidateKey(
             break;
 
         case AIRPDCAP_KEY_TYPE_WEP_40:
-            /* set the standard length and use a generic WEP key type  */
+            /* set the standard length and use a generic WEP key type */
             key->KeyData.Wep.WepKeyLen=AIRPDCAP_WEP_40_KEY_LEN;
             key->KeyType=AIRPDCAP_KEY_TYPE_WEP;
             break;
 
         case AIRPDCAP_KEY_TYPE_WEP_104:
-            /* set the standard length and use a generic WEP key type  */
+            /* set the standard length and use a generic WEP key type */
             key->KeyData.Wep.WepKeyLen=AIRPDCAP_WEP_104_KEY_LEN;
             key->KeyType=AIRPDCAP_KEY_TYPE_WEP;
             break;
 
         case AIRPDCAP_KEY_TYPE_WPA_PWD:
-            /* check passphrase and SSID size limits   */
+            /* check passphrase and SSID size limits */
             len=strlen(key->UserPwd.Passphrase);
             if (len<AIRPDCAP_WPA_PASSPHRASE_MIN_LEN || len>AIRPDCAP_WPA_PASSPHRASE_MAX_LEN) {
                 AIRPDCAP_DEBUG_PRINT_LINE("AirPDcapValidateKey", "WPA-PWD key: passphrase length not accepted", AIRPDCAP_DEBUG_LEVEL_5);
@@ -1097,16 +1510,15 @@ AirPDcapGetSa(
     PAIRPDCAP_CONTEXT ctx,
     AIRPDCAP_SEC_ASSOCIATION_ID *id)
 {
-    INT index;
-
-    if (ctx->last_stored_index!=-1) {
-        /* at least one association was stored                                                                                                         */
-        /* search for the association from last_stored_index to 0 (most recent added)  */
-        for (index=ctx->last_stored_index; index>=0; index--) {
-            if (ctx->sa[index].used) {
-                if (memcmp(id, &(ctx->sa[index].saId), sizeof(AIRPDCAP_SEC_ASSOCIATION_ID))==0) {
-                    ctx->index=index;
-                    return index;
+    INT sa_index;
+    if (ctx->sa_index!=-1) {
+        /* at least one association was stored                               */
+        /* search for the association from sa_index to 0 (most recent added) */
+        for (sa_index=ctx->sa_index; sa_index>=0; sa_index--) {
+            if (ctx->sa[sa_index].used) {
+                if (memcmp(id, &(ctx->sa[sa_index].saId), sizeof(AIRPDCAP_SEC_ASSOCIATION_ID))==0) {
+                    ctx->index=sa_index;
+                    return sa_index;
                 }
             }
         }
@@ -1121,41 +1533,44 @@ AirPDcapStoreSa(
     AIRPDCAP_SEC_ASSOCIATION_ID *id)
 {
     INT last_free;
-
+    if (ctx->first_free_index>=AIRPDCAP_MAX_SEC_ASSOCIATIONS_NR) {
+        /* there is no empty space available. FAILURE */
+        return -1;
+    }
     if (ctx->sa[ctx->first_free_index].used) {
-        /* last addition was in the middle of the array (and the first_free_index was just incremented by 1)   */
-        /* search for a free space from the first_free_index to AIRPDCAP_STA_INFOS_NR (to avoid free blocks in */
-        /*             the middle)                                                                                                                                                                                                                                     */
+        /* last addition was in the middle of the array (and the first_free_index was just incremented by 1)   */
+        /* search for a free space from the first_free_index to AIRPDCAP_STA_INFOS_NR (to avoid free blocks in */
+        /*              the middle)                                                                            */
         for (last_free=ctx->first_free_index; last_free<AIRPDCAP_MAX_SEC_ASSOCIATIONS_NR; last_free++)
             if (!ctx->sa[last_free].used)
                 break;
 
         if (last_free>=AIRPDCAP_MAX_SEC_ASSOCIATIONS_NR) {
-            /* there is no empty space available. FAILURE      */
+            /* there is no empty space available. FAILURE */
             return -1;
         }
 
-        /* store first free space index        */
+        /* store first free space index */
         ctx->first_free_index=last_free;
     }
 
-    /* use this info   */
+    /* use this info */
     ctx->index=ctx->first_free_index;
 
-    /* reset the info structure        */
+    /* reset the info structure */
     memset(ctx->sa+ctx->index, 0, sizeof(AIRPDCAP_SEC_ASSOCIATION));
 
     ctx->sa[ctx->index].used=1;
 
-    /* set the info structure  */
+    /* set the info structure */
     memcpy(&(ctx->sa[ctx->index].saId), id, sizeof(AIRPDCAP_SEC_ASSOCIATION_ID));
 
-    /* increment by 1 the first_free_index (heuristic) */
+    /* increment by 1 the first_free_index (heuristic) */
     ctx->first_free_index++;
 
-    /* set the last_stored_index if the added index is greater the the last_stored_index       */
-    if (ctx->index > ctx->last_stored_index)
-        ctx->last_stored_index=ctx->index;
+    /* set the sa_index if the added index is greater the the sa_index */
+    if (ctx->index > ctx->sa_index)
+        ctx->sa_index=ctx->index;
 
     return ctx->index;
 }
@@ -1179,10 +1594,15 @@ AirPDcapGetStaAddress(
     switch(AIRPDCAP_DS_BITS(frame->fc[1])) { /* Bit 1 = FromDS, bit 0 = ToDS */
         case 0:
         case 1:
-        case 3:
             return frame->addr2;
         case 2:
             return frame->addr1;
+        case 3:
+            if (memcmp(frame->addr1, frame->addr2, AIRPDCAP_MAC_LEN) < 0)
+                return frame->addr1;
+            else
+                return frame->addr2;
+
         default:
             return NULL;
     }
@@ -1196,33 +1616,40 @@ AirPDcapGetBssidAddress(
         case 0:
             return frame->addr3;
         case 1:
-        case 3:
             return frame->addr1;
         case 2:
             return frame->addr2;
+        case 3:
+            if (memcmp(frame->addr1, frame->addr2, AIRPDCAP_MAC_LEN) > 0)
+                return frame->addr1;
+            else
+                return frame->addr2;
+
         default:
             return NULL;
     }
 }
 
-/* Function used to derive the PTK. Refer to IEEE 802.11I-2004, pag. 74        */
+/* Function used to derive the PTK. Refer to IEEE 802.11I-2004, pag. 74
+ * and IEEE 802.11i-2004, pag. 164 */
 static void
 AirPDcapRsnaPrfX(
     AIRPDCAP_SEC_ASSOCIATION *sa,
     const UCHAR pmk[32],
     const UCHAR snonce[32],
-    const INT x,       /*      for TKIP 512, for CCMP 384      */
+    const INT x,        /*      for TKIP 512, for CCMP 384 */
     UCHAR *ptk)
 {
     UINT8 i;
     UCHAR R[100];
     INT offset=sizeof("Pairwise key expansion");
+    UCHAR output[80]; /* allow for sha1 overflow. */
 
     memset(R, 0, 100);
 
     memcpy(R, "Pairwise key expansion", offset);
 
-    /* Min(AA, SPA) || Max(AA, SPA)    */
+    /* Min(AA, SPA) || Max(AA, SPA) */
     if (memcmp(sa->saId.sta, sa->saId.bssid, AIRPDCAP_MAC_LEN) < 0)
     {
         memcpy(R + offset, sa->saId.sta, AIRPDCAP_MAC_LEN);
@@ -1236,7 +1663,7 @@ AirPDcapRsnaPrfX(
 
     offset+=AIRPDCAP_MAC_LEN*2;
 
-    /* Min(ANonce,SNonce) || Max(ANonce,SNonce)        */
+    /* Min(ANonce,SNonce) || Max(ANonce,SNonce) */
     if( memcmp(snonce, sa->wpa.nonce, 32) < 0 )
     {
         memcpy(R + offset, snonce, 32);
@@ -1253,10 +1680,13 @@ AirPDcapRsnaPrfX(
     for(i = 0; i < (x+159)/160; i++)
     {
         R[offset] = i;
-        sha1_hmac(pmk, 32, R, 100, ptk + i * 20);
+        sha1_hmac(pmk, 32, R, 100, &output[20 * i]);
     }
+    memcpy(ptk, output, x/8);
 }
 
+#define MAX_SSID_LENGTH 32 /* maximum SSID length */
+
 static INT
 AirPDcapRsnaPwd2PskStep(
     const guint8 *ppBytes,
@@ -1267,26 +1697,35 @@ AirPDcapRsnaPwd2PskStep(
     const INT count,
     UCHAR *output)
 {
-    UCHAR digest[36], digest1[AIRPDCAP_SHA_DIGEST_LEN];
+    UCHAR digest[MAX_SSID_LENGTH+4];  /* SSID plus 4 bytes of count */
+    UCHAR digest1[SHA1_DIGEST_LEN];
     INT i, j;
 
+    if (ssidLength > MAX_SSID_LENGTH) {
+        /* This "should not happen" */
+        return AIRPDCAP_RET_UNSUCCESS;
+    }
+
+    memset(digest, 0, sizeof digest);
+    memset(digest1, 0, sizeof digest1);
+
     /* U1 = PRF(P, S || INT(i)) */
     memcpy(digest, ssid, ssidLength);
     digest[ssidLength] = (UCHAR)((count>>24) & 0xff);
     digest[ssidLength+1] = (UCHAR)((count>>16) & 0xff);
     digest[ssidLength+2] = (UCHAR)((count>>8) & 0xff);
     digest[ssidLength+3] = (UCHAR)(count & 0xff);
-    sha1_hmac(ppBytes, ppLength, digest, ssidLength+4, digest1);
+    sha1_hmac(ppBytes, ppLength, digest, (guint32) ssidLength+4, digest1);
 
     /* output = U1 */
-    memcpy(output, digest1, AIRPDCAP_SHA_DIGEST_LEN);
+    memcpy(output, digest1, SHA1_DIGEST_LEN);
     for (i = 1; i < iterations; i++) {
         /* Un = PRF(P, Un-1) */
-        sha1_hmac(ppBytes, ppLength, digest1, AIRPDCAP_SHA_DIGEST_LEN, digest);
+        sha1_hmac(ppBytes, ppLength, digest1, SHA1_DIGEST_LEN, digest);
 
-        memcpy(digest1, digest, AIRPDCAP_SHA_DIGEST_LEN);
+        memcpy(digest1, digest, SHA1_DIGEST_LEN);
         /* output = output xor Un */
-        for (j = 0; j < AIRPDCAP_SHA_DIGEST_LEN; j++) {
+        for (j = 0; j < SHA1_DIGEST_LEN; j++) {
             output[j] ^= digest[j];
         }
     }
@@ -1301,10 +1740,10 @@ AirPDcapRsnaPwd2Psk(
     const size_t ssidLength,
     UCHAR *output)
 {
-    UCHAR m_output[AIRPDCAP_WPA_PSK_LEN];
+    UCHAR m_output[2*SHA1_DIGEST_LEN];
     GByteArray *pp_ba = g_byte_array_new();
 
-    memset(m_output, 0, AIRPDCAP_WPA_PSK_LEN);
+    memset(m_output, 0, 2*SHA1_DIGEST_LEN);
 
     if (!uri_str_to_bytes(passphrase, pp_ba)) {
         g_byte_array_free(pp_ba, TRUE);
@@ -1312,7 +1751,7 @@ AirPDcapRsnaPwd2Psk(
     }
 
     AirPDcapRsnaPwd2PskStep(pp_ba->data, pp_ba->len, ssid, ssidLength, 4096, 1, m_output);
-    AirPDcapRsnaPwd2PskStep(pp_ba->data, pp_ba->len, ssid, ssidLength, 4096, 2, &m_output[AIRPDCAP_SHA_DIGEST_LEN]);
+    AirPDcapRsnaPwd2PskStep(pp_ba->data, pp_ba->len, ssid, ssidLength, 4096, 2, &m_output[SHA1_DIGEST_LEN]);
 
     memcpy(output, m_output, AIRPDCAP_WPA_PSK_LEN);
     g_byte_array_free(pp_ba, TRUE);
@@ -1322,13 +1761,12 @@ AirPDcapRsnaPwd2Psk(
 
 /*
  * Returns the decryption_key_t struct given a string describing the key.
- * Returns NULL if the key_string cannot be parsed.
+ * Returns NULL if the input_string cannot be parsed.
  */
 decryption_key_t*
-parse_key_string(gchar* input_string)
+parse_key_string(gchar* input_string, guint8 key_type)
 {
-    gchar *type;
-    gchar *key;
+    gchar *key, *tmp_str;
     gchar *ssid;
 
     GString    *key_string = NULL;
@@ -1338,126 +1776,86 @@ parse_key_string(gchar* input_string)
     gchar **tokens;
     guint n = 0;
     decryption_key_t *dk;
-    gchar *first_nibble = input_string;
 
     if(input_string == NULL)
         return NULL;
 
     /*
-     * Parse the input_string. It should be in the form
-     * <key type>:<key data>[:<ssid>]
-     * XXX - For backward compatibility, the a WEP key can be just a string
-     * of hexadecimal characters (if WEP key is wrong, null will be
+     * Parse the input_string. WEP and WPA will be just a string
+     * of hexadecimal characters (if key is wrong, null will be
      * returned...).
+     * WPA-PWD should be in the form
+     * <key data>[:<ssid>]
      */
 
-    /* First, check for a WEP string */
-    /* XXX - This duplicates code in packet-ieee80211.c */
-    if (g_strncasecmp(input_string, STRING_KEY_TYPE_WEP ":", 4) == 0) {
-        first_nibble += 4;
-    }
-
-    key_ba = g_byte_array_new();
-    res = hex_str_to_bytes(first_nibble, key_ba, FALSE);
-
-    if (res && key_ba->len > 0) {
-        /* Key is correct! It was probably an 'old style' WEP key */
-        /* Create the decryption_key_t structure, fill it and return it*/
-        dk = g_malloc(sizeof(decryption_key_t));
-
-        dk->type = AIRPDCAP_KEY_TYPE_WEP;
-        /* XXX - The current key handling code in the GUI requires
-         * no separators and lower case */
-        dk->key  = g_string_new(bytes_to_str(key_ba->data, key_ba->len));
-        g_string_down(dk->key);
-        dk->bits = key_ba->len * 8;
-        dk->ssid = NULL;
-
-        g_byte_array_free(key_ba, TRUE);
-        return dk;
-    }
-    g_byte_array_free(key_ba, TRUE);
-
+    switch(key_type)
+    {
+    case AIRPDCAP_KEY_TYPE_WEP:
+    case AIRPDCAP_KEY_TYPE_WEP_40:
+    case AIRPDCAP_KEY_TYPE_WEP_104:
 
-    tokens = g_strsplit(input_string,":",0);
+       key_ba = g_byte_array_new();
+       res = hex_str_to_bytes(input_string, key_ba, FALSE);
 
-    /* Tokens is a null termiated array of strings ... */
-    while(tokens[n] != NULL)
-        n++;
+       if (res && key_ba->len > 0) {
+           /* Key is correct! It was probably an 'old style' WEP key */
+           /* Create the decryption_key_t structure, fill it and return it*/
+           dk = (decryption_key_t *)g_malloc(sizeof(decryption_key_t));
 
-    if(n < 2)
-    {
-        /* Free the array of strings */
-        g_strfreev(tokens);
-        return NULL;
-    }
+           dk->type = AIRPDCAP_KEY_TYPE_WEP;
+           /* XXX - The current key handling code in the GUI requires
+            * no separators and lower case */
+           tmp_str = bytes_to_str(NULL, key_ba->data, key_ba->len);
+           dk->key  = g_string_new(tmp_str);
+           g_string_ascii_down(dk->key);
+           dk->bits = key_ba->len * 8;
+           dk->ssid = NULL;
 
-    type = g_strdup(tokens[0]);
+           wmem_free(NULL, tmp_str);
+           g_byte_array_free(key_ba, TRUE);
+           return dk;
+       }
 
-    /*
-     * The second token is the key (right now it doesn't matter
-     * if it is a passphrase[+ssid] or an hexadecimal one)
-     */
-    key = g_strdup(tokens[1]);
+       /* Key doesn't work */
+       g_byte_array_free(key_ba, TRUE);
+       return NULL;
 
-    ssid = NULL;
-    /* Maybe there is a third token (an ssid, if everything else is ok) */
-    if(n >= 3)
-    {
-        ssid = g_strdup(tokens[2]);
-    }
+    case AIRPDCAP_KEY_TYPE_WPA_PWD:
 
-    if (g_strcasecmp(type,STRING_KEY_TYPE_WPA_PSK) == 0) /* WPA key */
-    {
-        /* Create a new string */
-        key_string = g_string_new(key);
+        tokens = g_strsplit(input_string,":",0);
 
-        key_ba = g_byte_array_new();
-        res = hex_str_to_bytes(key, key_ba, FALSE);
+        /* Tokens is a null termiated array of strings ... */
+        while(tokens[n] != NULL)
+            n++;
 
-        /* Two tokens means that the user should have entered a WPA-BIN key ... */
-        if(!res || ((key_string->len) != WPA_PSK_KEY_CHAR_SIZE))
+        if(n < 1)
         {
-            g_string_free(key_string, TRUE);
-            g_byte_array_free(key_ba, TRUE);
-
-            g_free(type);
-            g_free(key);
-            /* No ssid has been created ... */
             /* Free the array of strings */
             g_strfreev(tokens);
             return NULL;
         }
 
-        /* Key was correct!!! Create the new decryption_key_t ... */
-        dk = (decryption_key_t*)g_malloc(sizeof(decryption_key_t));
-
-        dk->type = AIRPDCAP_KEY_TYPE_WPA_PMK;
-        dk->key  = g_string_new(key);
-        dk->bits = dk->key->len * 4;
-        dk->ssid = NULL;
+        /*
+         * The first token is the key
+         */
+        key = g_strdup(tokens[0]);
 
-        g_string_free(key_string, TRUE);
-        g_byte_array_free(key_ba, TRUE);
-        g_free(key);
-        g_free(type);
+        ssid = NULL;
+        /* Maybe there is a second token (an ssid, if everything else is ok) */
+        if(n >= 2)
+        {
+           ssid = g_strdup(tokens[1]);
+        }
 
-        /* Free the array of strings */
-        g_strfreev(tokens);
-        return dk;
-    }
-    else if(g_strcasecmp(type,STRING_KEY_TYPE_WPA_PWD) == 0) /* WPA key *//* If the number of tokens is more than three, we accept the string... if the first three tokens are correct... */
-    {
         /* Create a new string */
         key_string = g_string_new(key);
         ssid_ba = NULL;
 
-        /* Three (or more) tokens mean that the user entered a WPA-PWD key ... */
+        /* Two (or more) tokens mean that the user entered a WPA-PWD key ... */
         if( ((key_string->len) > WPA_KEY_MAX_CHAR_SIZE) || ((key_string->len) < WPA_KEY_MIN_CHAR_SIZE))
         {
             g_string_free(key_string, TRUE);
 
-            g_free(type);
             g_free(key);
             g_free(ssid);
 
@@ -1466,13 +1864,12 @@ parse_key_string(gchar* input_string)
             return NULL;
         }
 
-        if(ssid != NULL) /* more than three tokens found, means that the user specified the ssid */
+        if(ssid != NULL) /* more than two tokens found, means that the user specified the ssid */
         {
             ssid_ba = g_byte_array_new();
             if (! uri_str_to_bytes(ssid, ssid_ba)) {
                 g_string_free(key_string, TRUE);
                 g_byte_array_free(ssid_ba, TRUE);
-                g_free(type);
                 g_free(key);
                 g_free(ssid);
                 /* Free the array of strings */
@@ -1485,7 +1882,6 @@ parse_key_string(gchar* input_string)
                 g_string_free(key_string, TRUE);
                 g_byte_array_free(ssid_ba, TRUE);
 
-                g_free(type);
                 g_free(key);
                 g_free(ssid);
 
@@ -1507,7 +1903,6 @@ parse_key_string(gchar* input_string)
         if (ssid_ba != NULL)
             g_byte_array_free(ssid_ba, TRUE);
 
-        g_free(type);
         g_free(key);
         if(ssid != NULL)
             g_free(ssid);
@@ -1515,23 +1910,47 @@ parse_key_string(gchar* input_string)
         /* Free the array of strings */
         g_strfreev(tokens);
         return dk;
-    }
 
-    /* Something was wrong ... free everything */
+    case AIRPDCAP_KEY_TYPE_WPA_PSK:
 
-    g_free(type);
-    g_free(key);
-    if(ssid != NULL)
-        g_free(ssid); /* It is not always present */
-    if (ssid_ba != NULL)
-        g_byte_array_free(ssid_ba, TRUE);
+        key_ba = g_byte_array_new();
+        res = hex_str_to_bytes(input_string, key_ba, FALSE);
 
-    /* Free the array of strings */
-    g_strfreev(tokens);
+        /* Two tokens means that the user should have entered a WPA-BIN key ... */
+        if(!res || ((key_ba->len) != WPA_PSK_KEY_SIZE))
+        {
+            g_byte_array_free(key_ba, TRUE);
+
+            /* No ssid has been created ... */
+            return NULL;
+        }
+
+        /* Key was correct!!! Create the new decryption_key_t ... */
+        dk = (decryption_key_t*)g_malloc(sizeof(decryption_key_t));
+
+        dk->type = AIRPDCAP_KEY_TYPE_WPA_PSK;
+        dk->key  = g_string_new(input_string);
+        dk->bits = (guint) dk->key->len * 4;
+        dk->ssid = NULL;
+
+        g_byte_array_free(key_ba, TRUE);
+        return dk;
+    }
 
+    /* Type not supported */
     return NULL;
 }
 
+void
+free_key_string(decryption_key_t *dk)
+{
+    if (dk->key)
+        g_string_free(dk->key, TRUE);
+    if (dk->ssid)
+        g_byte_array_free(dk->ssid, TRUE);
+    g_free(dk);
+}
+
 /*
  * Returns a newly allocated string representing the given decryption_key_t
  * struct, or NULL if something is wrong...
@@ -1546,28 +1965,26 @@ get_key_string(decryption_key_t* dk)
 
     switch(dk->type) {
         case AIRPDCAP_KEY_TYPE_WEP:
-            output_string = g_strdup_printf("%s:%s",STRING_KEY_TYPE_WEP,dk->key->str);
+            output_string = g_strdup(dk->key->str);
             break;
         case AIRPDCAP_KEY_TYPE_WPA_PWD:
             if(dk->ssid == NULL)
-                output_string = g_strdup_printf("%s:%s",STRING_KEY_TYPE_WPA_PWD,dk->key->str);
+                output_string = g_strdup(dk->key->str);
             else
-                output_string = g_strdup_printf("%s:%s:%s",
-                    STRING_KEY_TYPE_WPA_PWD, dk->key->str,
-                    format_uri(dk->ssid, ":"));
+                output_string = g_strdup_printf("%s:%s",
+                    dk->key->str, format_uri(dk->ssid, ":"));
             break;
         case AIRPDCAP_KEY_TYPE_WPA_PMK:
-            output_string = g_strdup_printf("%s:%s",STRING_KEY_TYPE_WPA_PSK,dk->key->str);
+            output_string = g_strdup(dk->key->str);
             break;
         default:
             return NULL;
-            break;
     }
 
     return output_string;
 }
 
-#ifdef __cplusplus
+#ifdef __cplusplus
 }
 #endif
 
@@ -1579,9 +1996,9 @@ get_key_string(decryption_key_t* dk)
  * Local Variables:
  * c-basic-offset: 4
  * tab-width: 8
- * set-tabs-mode: nil
+ * indent-tabs-mode: nil
  * End:
  *
- * ex: set shiftwidth=4 tabstop=8 expandtab
+ * ex: set shiftwidth=4 tabstop=8 expandtab:
  * :indentSize=4:tabSize=8:noTabs=true:
  */