/*
  * Copyright (c) 2016 MediaTek Inc.  All rights reserved.
  *
  * This software is available to you under a choice of one of two
  * licenses.  You may choose to be licensed under the terms of the GNU
  * General Public License (GPL) Version 2, available from the file
  * COPYING in the main directory of this source tree, or the
  * BSD license below:
  *
  *     Redistribution and use in source and binary forms, with or
  *     without modification, are permitted provided that the following
  *     conditions are met:
  *
  *      - Redistributions of source code must retain the above
  *        copyright notice, this list of conditions and the following
  *        disclaimer.
  *
  *      - Redistributions in binary form must reproduce the above
  *        copyright notice, this list of conditions and the following
  *        disclaimer in the documentation and/or other materials
  *        provided with the distribution.
  *
  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
  * BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
  * ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
  * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  * SOFTWARE.
  */


#include	"rt_config.h"

/*****************************/
/******** SBOX Table *********/
/*****************************/

UCHAR SboxTable[256] = {
	0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5,
	0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76,
	0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0,
	0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0,
	0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc,
	0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15,
	0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a,
	0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75,
	0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0,
	0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84,
	0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b,
	0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf,
	0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85,
	0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8,
	0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5,
	0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2,
	0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17,
	0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73,
	0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88,
	0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb,
	0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c,
	0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79,
	0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9,
	0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08,
	0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6,
	0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a,
	0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e,
	0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e,
	0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94,
	0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf,
	0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68,
	0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16
};

static VOID xor_32(IN PUCHAR a, IN PUCHAR b, OUT PUCHAR out)
{
	INT i;

	for (i = 0; i < 4; i++)
		out[i] = a[i] ^ b[i];
}

static VOID xor_128(IN PUCHAR a, IN PUCHAR b, OUT PUCHAR out)
{
	INT i;

	for (i = 0; i < 16; i++)
		out[i] = a[i] ^ b[i];
}

static UCHAR RTMPCkipSbox(IN UCHAR a)
{
	return SboxTable[(int)a];
}

static VOID next_key(IN PUCHAR key, IN INT round)
{
	UCHAR rcon;
	UCHAR sbox_key[4];
	UCHAR rcon_table[12] = {
		0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80,
		0x1b, 0x36, 0x36, 0x36
	};

	sbox_key[0] = RTMPCkipSbox(key[13]);
	sbox_key[1] = RTMPCkipSbox(key[14]);
	sbox_key[2] = RTMPCkipSbox(key[15]);
	sbox_key[3] = RTMPCkipSbox(key[12]);

	rcon = rcon_table[round];

	xor_32(&key[0], sbox_key, &key[0]);
	key[0] = key[0] ^ rcon;

	xor_32(&key[4], &key[0], &key[4]);
	xor_32(&key[8], &key[4], &key[8]);
	xor_32(&key[12], &key[8], &key[12]);
}

static VOID byte_sub(IN PUCHAR in, OUT PUCHAR out)
{
	INT i;

	for (i = 0; i < 16; i++)
		out[i] = RTMPCkipSbox(in[i]);
}

/************************************/
/* bitwise_xor()                    */
/* A 128 bit, bitwise exclusive or  */
/************************************/

static void bitwise_xor(unsigned char *ina, unsigned char *inb, unsigned char *out)
{
	int i;

	for (i = 0; i < 16; i++)
		out[i] = ina[i] ^ inb[i];
}

static VOID shift_row(IN PUCHAR in, OUT PUCHAR out)
{
	out[0] = in[0];
	out[1] = in[5];
	out[2] = in[10];
	out[3] = in[15];
	out[4] = in[4];
	out[5] = in[9];
	out[6] = in[14];
	out[7] = in[3];
	out[8] = in[8];
	out[9] = in[13];
	out[10] = in[2];
	out[11] = in[7];
	out[12] = in[12];
	out[13] = in[1];
	out[14] = in[6];
	out[15] = in[11];
}

static VOID mix_column(IN PUCHAR in, OUT PUCHAR out)
{
	INT i;
	UCHAR add1b[4];
	UCHAR add1bf7[4];
	UCHAR rotl[4];
	UCHAR swap_halfs[4];
	UCHAR andf7[4];
	UCHAR rotr[4];
	UCHAR temp[4];
	UCHAR tempb[4];

	for (i = 0; i < 4; i++) {
		if ((in[i] & 0x80) == 0x80)
			add1b[i] = 0x1b;
		else
			add1b[i] = 0x00;
	}

	swap_halfs[0] = in[2];	/* Swap halfs */
	swap_halfs[1] = in[3];
	swap_halfs[2] = in[0];
	swap_halfs[3] = in[1];

	rotl[0] = in[3];	/* Rotate left 8 bits */
	rotl[1] = in[0];
	rotl[2] = in[1];
	rotl[3] = in[2];

	andf7[0] = in[0] & 0x7f;
	andf7[1] = in[1] & 0x7f;
	andf7[2] = in[2] & 0x7f;
	andf7[3] = in[3] & 0x7f;

	for (i = 3; i > 0; i--) {	/* logical shift left 1 bit */
		andf7[i] = andf7[i] << 1;
		if ((andf7[i - 1] & 0x80) == 0x80)
			andf7[i] = (andf7[i] | 0x01);
	}
	andf7[0] = andf7[0] << 1;
	andf7[0] = andf7[0] & 0xfe;

	xor_32(add1b, andf7, add1bf7);

	xor_32(in, add1bf7, rotr);

	temp[0] = rotr[0];	/* Rotate right 8 bits */
	rotr[0] = rotr[1];
	rotr[1] = rotr[2];
	rotr[2] = rotr[3];
	rotr[3] = temp[0];

	xor_32(add1bf7, rotr, temp);
	xor_32(swap_halfs, rotl, tempb);
	xor_32(temp, tempb, out);
}

/************************************************/
/* construct_mic_header1()                      */
/* Builds the first MIC header block from       */
/* header fields.                               */
/************************************************/

static void construct_mic_header1(unsigned char *mic_header1,
				  int header_length, unsigned char *mpdu)
{
	mic_header1[0] = (unsigned char)((header_length - 2) / 256);
	mic_header1[1] = (unsigned char)((header_length - 2) % 256);
	mic_header1[2] = mpdu[0] & 0xcf;	/* Mute CF poll & CF ack bits */
	mic_header1[3] = mpdu[1] & 0xc7;	/* Mute retry, more data and pwr mgt bits */
	mic_header1[4] = mpdu[4];	/* A1 */
	mic_header1[5] = mpdu[5];
	mic_header1[6] = mpdu[6];
	mic_header1[7] = mpdu[7];
	mic_header1[8] = mpdu[8];
	mic_header1[9] = mpdu[9];
	mic_header1[10] = mpdu[10];	/* A2 */
	mic_header1[11] = mpdu[11];
	mic_header1[12] = mpdu[12];
	mic_header1[13] = mpdu[13];
	mic_header1[14] = mpdu[14];
	mic_header1[15] = mpdu[15];
}

/************************************************/
/* construct_mic_header2()                      */
/* Builds the last MIC header block from        */
/* header fields.                               */
/************************************************/

static void construct_mic_header2(unsigned char *mic_header2,
				  unsigned char *mpdu, int a4_exists, int qc_exists)
{
	int i;

	for (i = 0; i < 16; i++)
		mic_header2[i] = 0x00;

	mic_header2[0] = mpdu[16];	/* A3 */
	mic_header2[1] = mpdu[17];
	mic_header2[2] = mpdu[18];
	mic_header2[3] = mpdu[19];
	mic_header2[4] = mpdu[20];
	mic_header2[5] = mpdu[21];

	/* In Sequence Control field, mute sequence numer bits (12-bit) */
	mic_header2[6] = mpdu[22] & 0x0f;	/* SC */
	mic_header2[7] = 0x00;	/* mpdu[23]; */

	if ((!qc_exists) & a4_exists) {
		for (i = 0; i < 6; i++)
			mic_header2[8 + i] = mpdu[24 + i];	/* A4 */

	}

	if (qc_exists && (!a4_exists)) {
		mic_header2[8] = mpdu[24] & 0x0f;	/* mute bits 15 - 4 */
		mic_header2[9] = mpdu[25] & 0x00;
	}

	if (qc_exists && a4_exists) {
		for (i = 0; i < 6; i++)
			mic_header2[8 + i] = mpdu[24 + i];	/* A4 */

		mic_header2[14] = mpdu[30] & 0x0f;
		mic_header2[15] = mpdu[31] & 0x00;
	}
}

/************************************************/
/* construct_mic_iv()                           */
/* Builds the MIC IV from header fields and PN  */
/************************************************/

static void construct_mic_iv(unsigned char *mic_iv,
			     int qc_exists,
			     int a4_exists,
			     unsigned char *mpdu,
			     unsigned int payload_length, unsigned char *pn_vector)
{
	int i;

	mic_iv[0] = 0x59;
	if (qc_exists && a4_exists)
		mic_iv[1] = mpdu[30] & 0x0f;	/* QoS_TC           */
	if (qc_exists && !a4_exists)
		mic_iv[1] = mpdu[24] & 0x0f;	/* mute bits 7-4    */
	if (!qc_exists)
		mic_iv[1] = 0x00;
	for (i = 2; i < 8; i++)
		mic_iv[i] = mpdu[i + 8];	/* mic_iv[2:7] = A2[0:5] = mpdu[10:15] */
#ifdef CONSISTENT_PN_ORDER
	for (i = 8; i < 14; i++)
		mic_iv[i] = pn_vector[i - 8];	/* mic_iv[8:13] = PN[0:5] */
#else
	for (i = 8; i < 14; i++)
		mic_iv[i] = pn_vector[13 - i];	/* mic_iv[8:13] = PN[5:0] */
#endif /* endif */
	mic_iv[14] = (unsigned char)(payload_length / 256);
	mic_iv[15] = (unsigned char)(payload_length % 256);

}

/****************************************/
/* aes128k128d()                        */
/* Performs a 128 bit AES encrypt with  */
/* 128 bit data.                        */
/****************************************/
static void aes128k128d(unsigned char *key, unsigned char *data, unsigned char *ciphertext)
{
	int round;
	int i;
	unsigned char intermediatea[16];
	unsigned char intermediateb[16];
	unsigned char round_key[16];

	for (i = 0; i < 16; i++)
		round_key[i] = key[i];

	for (round = 0; round < 11; round++) {
		if (round == 0) {
			xor_128(round_key, data, ciphertext);
			next_key(round_key, round);
		} else if (round == 10) {
			byte_sub(ciphertext, intermediatea);
			shift_row(intermediatea, intermediateb);
			xor_128(intermediateb, round_key, ciphertext);
		} else {	/* 1 - 9 */

			byte_sub(ciphertext, intermediatea);
			shift_row(intermediatea, intermediateb);
			mix_column(&intermediateb[0], &intermediatea[0]);
			mix_column(&intermediateb[4], &intermediatea[4]);
			mix_column(&intermediateb[8], &intermediatea[8]);
			mix_column(&intermediateb[12], &intermediatea[12]);
			xor_128(intermediatea, round_key, ciphertext);
			next_key(round_key, round);
		}
	}

}

static void construct_ctr_preload(unsigned char *ctr_preload,
				  int a4_exists,
				  int qc_exists,
				  unsigned char *mpdu, unsigned char *pn_vector, int c)
{

	int i = 0;

	for (i = 0; i < 16; i++)
		ctr_preload[i] = 0x00;
	i = 0;

	ctr_preload[0] = 0x01;	/* flag */
	if (qc_exists && a4_exists)
		ctr_preload[1] = mpdu[30] & 0x0f;	/* QoC_Control  */
	if (qc_exists && !a4_exists)
		ctr_preload[1] = mpdu[24] & 0x0f;

	for (i = 2; i < 8; i++)
		ctr_preload[i] = mpdu[i + 8];	/* ctr_preload[2:7] = A2[0:5] = mpdu[10:15] */
#ifdef CONSISTENT_PN_ORDER
	for (i = 8; i < 14; i++)
		ctr_preload[i] = pn_vector[i - 8];	/* ctr_preload[8:13] = PN[0:5] */
#else
	for (i = 8; i < 14; i++)
		ctr_preload[i] = pn_vector[13 - i];	/* ctr_preload[8:13] = PN[5:0] */
#endif /* endif */
	ctr_preload[14] = (unsigned char)(c / 256);	/* Ctr */
	ctr_preload[15] = (unsigned char)(c % 256);

}

BOOLEAN RTMPSoftDecryptAES(IN PRTMP_ADAPTER pAd,
			   IN PUCHAR pData, IN ULONG DataByteCnt, IN PCIPHER_KEY pWpaKey)
{
	UINT HeaderLen;
	UCHAR PN[6];
	UINT payload_len;
	UINT num_blocks;
	UINT payload_remainder;
	USHORT fc;
	UCHAR fc0;
	UCHAR fc1;
	UINT frame_type;
	UINT frame_subtype;
	UINT from_ds;
	UINT to_ds;
	INT a4_exists;
	INT qc_exists;
	UCHAR aes_out[16];
	int payload_index;
	UINT i;
	UCHAR ctr_preload[16];
	UCHAR chain_buffer[16];
	UCHAR padded_buffer[16];
	UCHAR mic_iv[16];
	UCHAR mic_header1[16];
	UCHAR mic_header2[16];
	UCHAR MIC[8];
	UCHAR TrailMIC[8];

#ifdef RT_BIG_ENDIAN
	RTMPFrameEndianChange(pAd, (PUCHAR) pData, DIR_READ, FALSE);
#endif /* endif */

	fc0 = *pData;
	fc1 = *(pData + 1);

	fc = *((PUSHORT) pData);

	frame_type = ((fc0 >> 2) & 0x03);
	frame_subtype = ((fc0 >> 4) & 0x0f);

	from_ds = (fc1 & 0x2) >> 1;
	to_ds = (fc1 & 0x1);

	a4_exists = (from_ds & to_ds);
	qc_exists = ((frame_subtype == 0x08) ||	/* Assumed QoS subtypes */
		     (frame_subtype == 0x09) ||	/* Likely to change.    */
		     (frame_subtype == 0x0a) || (frame_subtype == 0x0b)
	    );

	HeaderLen = 24;

	if (a4_exists)
		HeaderLen += 6;

	if (qc_exists)
		HeaderLen += 2;

	if (pWpaKey->KeyLen == 0) {
		DBGPRINT(RT_DEBUG_TRACE, ("RTMPSoftDecryptAES failed!(the Length can not be 0)\n"));
		return FALSE;
	}

	PN[0] = *(pData + HeaderLen);
	PN[1] = *(pData + HeaderLen + 1);
	PN[2] = *(pData + HeaderLen + 4);
	PN[3] = *(pData + HeaderLen + 5);
	PN[4] = *(pData + HeaderLen + 6);
	PN[5] = *(pData + HeaderLen + 7);

	payload_len = DataByteCnt - HeaderLen - 8 - 8;	/* 8 bytes for CCMP header , 8 bytes for MIC */
	payload_remainder = (payload_len) % 16;
	num_blocks = (payload_len) / 16;


	/* Find start of payload */
	payload_index = HeaderLen + 8;	/*IV+EIV */

	for (i = 0; i < num_blocks; i++) {
		construct_ctr_preload(ctr_preload, a4_exists, qc_exists, pData, PN, i + 1);

		aes128k128d(pWpaKey->Key, ctr_preload, aes_out);

		bitwise_xor(aes_out, pData + payload_index, chain_buffer);
		NdisMoveMemory(pData + payload_index - 8, chain_buffer, 16);
		payload_index += 16;
	}

	/* If there is a short final block, then pad it */
	/* encrypt it and copy the unpadded part back */

	if (payload_remainder > 0) {
		construct_ctr_preload(ctr_preload, a4_exists, qc_exists, pData, PN, num_blocks + 1);

		NdisZeroMemory(padded_buffer, 16);
		NdisMoveMemory(padded_buffer, pData + payload_index, payload_remainder);

		aes128k128d(pWpaKey->Key, ctr_preload, aes_out);

		bitwise_xor(aes_out, padded_buffer, chain_buffer);
		NdisMoveMemory(pData + payload_index - 8, chain_buffer, payload_remainder);
		payload_index += payload_remainder;
	}

	/* Descrypt the MIC */
	construct_ctr_preload(ctr_preload, a4_exists, qc_exists, pData, PN, 0);
	NdisZeroMemory(padded_buffer, 16);
	NdisMoveMemory(padded_buffer, pData + payload_index, 8);

	aes128k128d(pWpaKey->Key, ctr_preload, aes_out);

	bitwise_xor(aes_out, padded_buffer, chain_buffer);

	NdisMoveMemory(TrailMIC, chain_buffer, 8);


	/* Calculate MIC */

	/*Force the protected frame bit on */
	*(pData + 1) = *(pData + 1) | 0x40;

	/* Find start of payload */
	/* Because the CCMP header has been removed */
	payload_index = HeaderLen;

	construct_mic_iv(mic_iv, qc_exists, a4_exists, pData, payload_len, PN);

	construct_mic_header1(mic_header1, HeaderLen, pData);

	construct_mic_header2(mic_header2, pData, a4_exists, qc_exists);

	aes128k128d(pWpaKey->Key, mic_iv, aes_out);
	bitwise_xor(aes_out, mic_header1, chain_buffer);
	aes128k128d(pWpaKey->Key, chain_buffer, aes_out);
	bitwise_xor(aes_out, mic_header2, chain_buffer);
	aes128k128d(pWpaKey->Key, chain_buffer, aes_out);

	/* iterate through each 16 byte payload block */
	for (i = 0; i < num_blocks; i++) {
		bitwise_xor(aes_out, pData + payload_index, chain_buffer);
		payload_index += 16;
		aes128k128d(pWpaKey->Key, chain_buffer, aes_out);
	}

	/* Add on the final payload block if it needs padding */
	if (payload_remainder > 0) {
		NdisZeroMemory(padded_buffer, 16);
		NdisMoveMemory(padded_buffer, pData + payload_index, payload_remainder);

		bitwise_xor(aes_out, padded_buffer, chain_buffer);
		aes128k128d(pWpaKey->Key, chain_buffer, aes_out);
	}
	/* aes_out contains padded mic, discard most significant */
	/* 8 bytes to generate 64 bit MIC */
	for (i = 0; i < 8; i++)
		MIC[i] = aes_out[i];

	if (!NdisEqualMemory(MIC, TrailMIC, 8)) {
		DBGPRINT(RT_DEBUG_ERROR, ("RTMPSoftDecryptAES, MIC Error !\n"));	/* MIC error. */
		return FALSE;
	}
#ifdef RT_BIG_ENDIAN
	RTMPFrameEndianChange(pAd, (PUCHAR) pData, DIR_READ, FALSE);
#endif /* endif */

	return TRUE;
}


/*
	========================================================================

	Routine Description:
		Construct AAD of CCMP.

	Arguments:

	Return Value:

	Note:
		It's described in IEEE Std 802.11-2007.
		The AAD is constructed from the MPDU header.

	========================================================================
*/
static VOID RTMPConstructCCMPAAD(IN PUCHAR pHdr,
				 IN BOOLEAN isDataFrame,
				 IN UINT8 a4_exists,
				 IN UINT8 qc_exists, OUT UCHAR *aad_hdr, OUT UINT *aad_len)
{
	UINT len = 0;

	/* Frame control -
	   Subtype bits (bits 4 5 6) in a Data MPDU masked to 0
	   Retry bit (bit 11) masked to 0
	   PwrMgt bit (bit 12) masked to 0
	   MoreData bit (bit 13) masked to 0
	   Protected Frame bit (bit 14) always set to 1 */
	if (isDataFrame)
		aad_hdr[0] = (*pHdr) & 0x8f;
	else
		aad_hdr[0] = (*pHdr);
	aad_hdr[1] = (*(pHdr + 1)) & 0xc7;
	aad_hdr[1] = aad_hdr[1] | 0x40;
	len = 2;

	/* Append Addr 1, 2 & 3 */
	NdisMoveMemory(&aad_hdr[len], pHdr + 4, 3 * MAC_ADDR_LEN);
	len += (3 * MAC_ADDR_LEN);

	/*  SC -
	   MPDU Sequence Control field, with the Sequence Number
	   subfield (bits 4-15 of the Sequence Control field)
	   masked to 0. The Fragment Number subfield is not modified. */
	aad_hdr[len] = (*(pHdr + 22)) & 0x0f;
	aad_hdr[len + 1] = 0x00;
	len += 2;

	/* Append the Addr4 field if present. */
	if (a4_exists) {
		NdisMoveMemory(&aad_hdr[len], pHdr + 24, MAC_ADDR_LEN);
		len += MAC_ADDR_LEN;
	}

	/*  QC -
	   QoS Control field, if present, a 2-octet field that includes
	   the MSDU priority. The QC TID field is used in the
	   construction of the AAD and the remaining QC fields are
	   set to 0 for the AAD calculation (bits 4 to 15 are set to 0). */
	if (qc_exists & a4_exists) {
		aad_hdr[len] = (*(pHdr + 30)) & 0x0f;	/* Qos_TC */
		aad_hdr[len + 1] = 0x00;
		len += 2;
	} else if (qc_exists & !a4_exists) {
		aad_hdr[len] = (*(pHdr + 24)) & 0x0f;	/* Qos_TC */
		aad_hdr[len + 1] = 0x00;
		len += 2;
	}

	*aad_len = len;
}

/*
	========================================================================

	Routine Description:
		Construct NONCE header of CCMP.

	Arguments:

	Return Value:

	Note:

	========================================================================
*/
static VOID RTMPConstructCCMPNonce(IN PUCHAR pHdr,
				   IN UINT8 a4_exists,
				   IN UINT8 qc_exists,
				   IN BOOLEAN isMgmtFrame,
				   IN UCHAR *pn, OUT UCHAR *nonce_hdr, OUT UINT *nonce_hdr_len)
{
	UINT n_offset = 0;
	INT i;

	/* Decide the Priority Octet
	   The Priority sub-field of the Nonce Flags field shall
	   be set to the fixed value 0 when there is no QC field
	   present in the MPDU header. When the QC field is present,
	   bits 0 to 3 of the Priority field shall be set to the
	   value of the QC TID (bits 0 to 3 of the QC field). */
	if (qc_exists && a4_exists)
		nonce_hdr[0] = (*(pHdr + 30)) & 0x0f;
	if (qc_exists && !a4_exists)
		nonce_hdr[0] = (*(pHdr + 24)) & 0x0f;

#ifdef DOT11W_PMF_SUPPORT
	/* When Management Frame Protection is negotiated, the Management
	   field of the Nonce Flags field shall be set to 1 if the Type
	   field of the Frame Control field is 00 (Management frame); otherwise it
	   is set to 0.  */
	if (isMgmtFrame)
		nonce_hdr[0] = nonce_hdr[0] | 0x10;
#endif /* DOT11W_PMF_SUPPORT */
	n_offset += 1;

	/* Fill in MPDU Address A2 field */
	NdisMoveMemory(&nonce_hdr[n_offset], pHdr + 10, MAC_ADDR_LEN);
	n_offset += MAC_ADDR_LEN;

	/* Fill in the PN. The PN field occupies octets 7¡V12.
	   The octets of PN shall be ordered so that PN0 is at octet index 12
	   and PN5 is at octet index 7. */
	for (i = 0; i < 6; i++)
		nonce_hdr[n_offset + i] = pn[5 - i];
	n_offset += LEN_PN;

	*nonce_hdr_len = n_offset;

}

/*
	========================================================================

	Routine Description:
		Construct CCMP header.

	Arguments:

	Return Value:

	Note:
		It's a 8-octets header.

	========================================================================
*/
VOID RTMPConstructCCMPHdr(IN UINT8 key_idx, IN UCHAR *pn, OUT UCHAR *ccmp_hdr)
{
	NdisZeroMemory(ccmp_hdr, LEN_CCMP_HDR);

	ccmp_hdr[0] = pn[0];
	ccmp_hdr[1] = pn[1];
	ccmp_hdr[3] = (key_idx << 6) | 0x20;
	ccmp_hdr[4] = pn[2];
	ccmp_hdr[5] = pn[3];
	ccmp_hdr[6] = pn[4];
	ccmp_hdr[7] = pn[5];
}

/*
	========================================================================

	Routine Description:

	Arguments:

	Return Value:

	Note:

	========================================================================
*/
BOOLEAN RTMPSoftEncryptCCMP(IN PRTMP_ADAPTER pAd,
			    IN PUCHAR pHdr,
			    IN PUCHAR pIV, IN PUCHAR pKey, INOUT PUCHAR pData, IN UINT32 DataLen)
{
	UINT8 frame_type, frame_subtype;
	UINT8 from_ds, to_ds;
	UINT8 a4_exists, qc_exists;
	UINT8 aad_hdr[30];
	UINT aad_len = 0;
	UINT8 nonce_hdr[13];
	UINT32 nonce_hdr_len = 0;
	UINT32 out_len = DataLen + 8;

#ifdef RT_BIG_ENDIAN
	RTMPFrameEndianChange(pAd, (PUCHAR) pHdr, DIR_READ, FALSE);
#endif /* endif */

	/* Initial variable */
	NdisZeroMemory(aad_hdr, 30);
	NdisZeroMemory(nonce_hdr, 13);

	/* Indicate type and subtype of Frame Control field */
	frame_type = (((*pHdr) >> 2) & 0x03);
	frame_subtype = (((*pHdr) >> 4) & 0x0f);

	/* Indicate the fromDS and ToDS */
	from_ds = ((*(pHdr + 1)) & 0x2) >> 1;
	to_ds = ((*(pHdr + 1)) & 0x1);

	/* decide if the Address 4 exist or QoS exist */
	a4_exists = (from_ds & to_ds);
	qc_exists = 0;
	if (frame_type == FC_TYPE_DATA) {
		qc_exists = ((frame_subtype == SUBTYPE_QDATA) ||
			     (frame_subtype == SUBTYPE_QDATA_CFACK) ||
			     (frame_subtype == SUBTYPE_QDATA_CFPOLL) ||
			     (frame_subtype == SUBTYPE_QDATA_CFACK_CFPOLL));
	}

	/* Construct AAD header */
	RTMPConstructCCMPAAD(pHdr,
			     (frame_type == FC_TYPE_DATA), a4_exists, qc_exists, aad_hdr, &aad_len);

	/* Construct NONCE header */
	RTMPConstructCCMPNonce(pHdr,
			       a4_exists,
			       qc_exists,
			       (frame_type == FC_TYPE_MGMT), pIV, nonce_hdr, &nonce_hdr_len);

	/* CCM originator processing -
	   Use the temporal key, AAD, nonce, and MPDU data to
	   form the cipher text and MIC. */
	if (AES_CCM_Encrypt(pData, DataLen,
			    pKey, 16,
			    nonce_hdr, nonce_hdr_len,
			    aad_hdr, aad_len, LEN_CCMP_MIC, pData, &out_len))
		return FALSE;

#ifdef RT_BIG_ENDIAN
	RTMPFrameEndianChange(pAd, (PUCHAR) pHdr, DIR_READ, FALSE);
#endif /* endif */

	return TRUE;
}

/*
	========================================================================

	Routine Description:
		Decrypt data with CCMP.

	Arguments:

	Return Value:

	Note:

	========================================================================
*/
BOOLEAN RTMPSoftDecryptCCMP(IN PRTMP_ADAPTER pAd,
			    IN PUCHAR pHdr,
			    IN PCIPHER_KEY pKey, INOUT PUCHAR pData, INOUT UINT16 * DataLen)
{
	UINT8 frame_type, frame_subtype;
	UINT8 from_ds, to_ds;
	UINT8 a4_exists, qc_exists;
	UINT8 aad_hdr[30];
	UINT aad_len = 0;
	UINT8 pn[LEN_PN];
	PUCHAR cipherData_ptr;
	UINT32 cipherData_len;
	UINT8 nonce_hdr[13];
	UINT32 nonce_hdr_len = 0;
	UINT32 out_len = *DataLen;

#ifdef RT_BIG_ENDIAN
	RTMPFrameEndianChange(pAd, (PUCHAR) pHdr, DIR_READ, FALSE);
#endif /* endif */

	/* Check the key is valid */
	if (pKey->KeyLen == 0) {
		DBGPRINT(RT_DEBUG_ERROR, ("%s : The key is not available !\n", __func__));
		return FALSE;
	}

	/* Initial variable */
	NdisZeroMemory(aad_hdr, 30);
	NdisZeroMemory(nonce_hdr, 13);

	/* Indicate type and subtype of Frame Control field */
	frame_type = (((*pHdr) >> 2) & 0x03);
	frame_subtype = (((*pHdr) >> 4) & 0x0f);

	/* Indicate the fromDS and ToDS */
	from_ds = ((*(pHdr + 1)) & 0x2) >> 1;
	to_ds = ((*(pHdr + 1)) & 0x1);

	/* decide if the Address 4 exist or QoS exist */
	a4_exists = (from_ds & to_ds);
	qc_exists = 0;
	if (frame_type == FC_TYPE_DATA) {
		qc_exists = ((frame_subtype == SUBTYPE_QDATA) ||
			     (frame_subtype == SUBTYPE_QDATA_CFACK) ||
			     (frame_subtype == SUBTYPE_QDATA_CFPOLL) ||
			     (frame_subtype == SUBTYPE_QDATA_CFACK_CFPOLL));
	}

	/* Extract PN and from CCMP header */
	pn[0] = pData[0];
	pn[1] = pData[1];
	pn[2] = pData[4];
	pn[3] = pData[5];
	pn[4] = pData[6];
	pn[5] = pData[7];

	/* skip ccmp header */
	cipherData_ptr = pData + LEN_CCMP_HDR;
	cipherData_len = *DataLen - LEN_CCMP_HDR;

	/*skip payload length is zero */
	if ((*DataLen) <= LEN_CCMP_HDR)
		return FALSE;

	/* Construct AAD header */
	RTMPConstructCCMPAAD(pHdr,
			     (frame_type == FC_TYPE_DATA), a4_exists, qc_exists, aad_hdr, &aad_len);

	/* Construct NONCE header */
	RTMPConstructCCMPNonce(pHdr,
			       a4_exists,
			       qc_exists,
			       (frame_type == FC_TYPE_MGMT), pn, nonce_hdr, &nonce_hdr_len);

	/* CCM recipient processing -
	   uses the temporal key, AAD, nonce, MIC,
	   and MPDU cipher text data */
	if (AES_CCM_Decrypt(cipherData_ptr, cipherData_len,
			    pKey->Key, 16,
			    nonce_hdr, nonce_hdr_len,
			    aad_hdr, aad_len, LEN_CCMP_MIC, pData, &out_len))
		return FALSE;

	*DataLen = out_len;

#ifdef RT_BIG_ENDIAN
	RTMPFrameEndianChange(pAd, (PUCHAR) pHdr, DIR_READ, FALSE);
#endif /* endif */

	return TRUE;
}

/*
	========================================================================

	Routine Description:
		CCMP test vector

	Arguments:

	Return Value:

	Note:

	========================================================================
*/
VOID CCMP_test_vector(IN PRTMP_ADAPTER pAd, IN INT input)
{
	UINT8 Key_ID = 0;
	/*UINT8 A1[6] =  {0x0f, 0xd2, 0xe1, 0x28, 0xa5, 0x7c}; */
	/*UINT8 A2[6] =  {0x50, 0x30, 0xf1, 0x84, 0x44, 0x08}; */
	/*UINT8 A3[6] =  {0xab, 0xae, 0xa5, 0xb8, 0xfc, 0xba}; */
	UINT8 TK[16] = { 0xc9, 0x7c, 0x1f, 0x67, 0xce, 0x37, 0x11, 0x85,
		0x51, 0x4a, 0x8a, 0x19, 0xf2, 0xbd, 0xd5, 0x2f
	};
	UINT8 PN[6] = { 0x0C, 0xE7, 0x76, 0x97, 0x03, 0xB5 };
	UINT8 HDR[24] = { 0x08, 0x48, 0xc3, 0x2c, 0x0f, 0xd2, 0xe1, 0x28,
		0xa5, 0x7c, 0x50, 0x30, 0xf1, 0x84, 0x44, 0x08,
		0xab, 0xae, 0xa5, 0xb8, 0xfc, 0xba, 0x80, 0x33
	};
	UINT8 AAD[22] = { 0x08, 0x40, 0x0f, 0xd2, 0xe1, 0x28, 0xa5, 0x7c,
		0x50, 0x30, 0xf1, 0x84, 0x44, 0x08, 0xab, 0xae,
		0xa5, 0xb8, 0xfc, 0xba, 0x00, 0x00
	};
	UINT8 CCMP_HDR[8] = { 0x0c, 0xe7, 0x00, 0x20, 0x76, 0x97, 0x03, 0xb5 };
	UINT8 CCM_NONCE[13] = { 0x00, 0x50, 0x30, 0xf1, 0x84, 0x44, 0x08, 0xb5,
		0x03, 0x97, 0x76, 0xe7, 0x0c
	};
	UINT8 P_TEXT_DATA[20] = { 0xf8, 0xba, 0x1a, 0x55, 0xd0, 0x2f, 0x85, 0xae,
		0x96, 0x7b, 0xb6, 0x2f, 0xb6, 0xcd, 0xa8, 0xeb,
		0x7e, 0x78, 0xa0, 0x50
	};
	UINT8 C_TEXT_DATA[28] = { 0xf3, 0xd0, 0xa2, 0xfe, 0x9a, 0x3d, 0xbf, 0x23,
		0x42, 0xa6, 0x43, 0xe4, 0x32, 0x46, 0xe8, 0x0c,
		0x3c, 0x04, 0xd0, 0x19, 0x78, 0x45, 0xce, 0x0b,
		0x16, 0xf9, 0x76, 0x23
	};
	UINT8 res_buf[100];
	UINT res_len = 0;

	DBGPRINT(RT_DEBUG_INFO, ("== CCMP test vector ==\n"));

	/* Check AAD */
	NdisZeroMemory(res_buf, 100);
	res_len = 0;
	RTMPConstructCCMPAAD(HDR, TRUE, 0, 0, res_buf, &res_len);
	if (res_len == 22 && NdisEqualMemory(res_buf, AAD, res_len))
		DBGPRINT(RT_DEBUG_INFO, ("Construct AAD is OK!!!\n"));
	else {
		DBGPRINT(RT_DEBUG_ERROR, ("\n!!!Construct AAD is FAILURE!!!\n\n"));
		hex_dump("Calculate AAD", res_buf, res_len);
	}
	/* Check NONCE */
	NdisZeroMemory(res_buf, 100);
	res_len = 0;
	RTMPConstructCCMPNonce(HDR, 0, 0, FALSE, PN, res_buf, &res_len);
	if (res_len == 13 && NdisEqualMemory(res_buf, CCM_NONCE, res_len))
		DBGPRINT(RT_DEBUG_INFO, ("Construct NONCE is OK!!!\n"));
	else {
		DBGPRINT(RT_DEBUG_ERROR, ("\n!!!Construct NONCE is FAILURE!!!\n\n"));
		hex_dump("Calculate NONCE", res_buf, res_len);
	}
	/* Check CCMP-Header */
	NdisZeroMemory(res_buf, 100);
	res_len = 0;
	RTMPConstructCCMPHdr(Key_ID, PN, res_buf);
	if (NdisEqualMemory(res_buf, CCMP_HDR, 8))
		DBGPRINT(RT_DEBUG_INFO, ("Construct CCMP_HDR is OK!!!\n"));
	else {
		DBGPRINT(RT_DEBUG_ERROR, ("\n!!!Construct CCMP_HDR is FAILURE!!!\n\n"));
		hex_dump("Calculate CCMP_HDR", res_buf, 8);
	}

	/* Encrypt action */
	NdisZeroMemory(res_buf, 100);
	NdisMoveMemory(res_buf, P_TEXT_DATA, sizeof(P_TEXT_DATA));
	res_len = sizeof(C_TEXT_DATA);
	if (AES_CCM_Encrypt(res_buf, sizeof(P_TEXT_DATA),
			    TK, sizeof(TK),
			    CCM_NONCE, sizeof(CCM_NONCE),
			    AAD, sizeof(AAD), 8, res_buf, &res_len) == 0) {
		if (res_len == sizeof(C_TEXT_DATA) &&
		    NdisEqualMemory(res_buf, C_TEXT_DATA, res_len))
			DBGPRINT(RT_DEBUG_INFO, ("CCM_Encrypt is OK!!!\n"));
		else {
			DBGPRINT(RT_DEBUG_ERROR, ("\n!!!CCM_Encrypt is FAILURE!!!\n\n"));
			hex_dump("CCM_Encrypt", res_buf, res_len);
		}
	}

	/* Decrypt action */
	NdisZeroMemory(res_buf, 100);
	NdisMoveMemory(res_buf, C_TEXT_DATA, sizeof(C_TEXT_DATA));
	res_len = sizeof(P_TEXT_DATA);
	if (AES_CCM_Decrypt(res_buf, sizeof(C_TEXT_DATA), TK, 16,
			    CCM_NONCE, sizeof(CCM_NONCE),
			    AAD, sizeof(AAD), 8, res_buf, &res_len) == 0) {
		if (res_len == sizeof(P_TEXT_DATA) &&
		    NdisEqualMemory(res_buf, P_TEXT_DATA, res_len))
			DBGPRINT(RT_DEBUG_INFO, ("CCM_Decrypt is OK!!!\n"));
		else {
			DBGPRINT(RT_DEBUG_ERROR, ("\n!!!CCM_Decrypt is FAILURE!!!\n\n"));
			hex_dump("CCM_Decrypt", res_buf, res_len);
		}
	}

	DBGPRINT(RT_DEBUG_INFO, ("== CCMP test vector ==\n"));

}

