/* rijndael-api-ref.c   v2.1   April 2000
 * Reference ANSI C code
 * authors: v2.0 Paulo Barreto
 *               Vincent Rijmen, K.U.Leuven
 *          v2.1 Vincent Rijmen, K.U.Leuven
 *
 * This code is placed in the public domain.
 */

#include "mvCommon.h"
#include "mvOs.h"

#include "mvAes.h"
#include "mvAesAlg.h"

/*  Defines:
	Add any additional defines you need
*/

#define     MODE_ECB        1	/*  Are we ciphering in ECB mode?   */
#define     MODE_CBC        2	/*  Are we ciphering in CBC mode?   */
#define     MODE_CFB1       3	/*  Are we ciphering in 1-bit CFB mode? */

int aesMakeKey(MV_U8 *expandedKey, MV_U8 *keyMaterial, int keyLen, int blockLen)
{
	MV_U8 W[MAXROUNDS + 1][4][MAXBC];
	MV_U8 k[4][MAXKC];
	MV_U8 j;
	int i, rounds, KC;

	if (expandedKey == NULL)
		return AES_BAD_KEY_INSTANCE;

	if (!((keyLen == 128) || (keyLen == 192) || (keyLen == 256)))
		return AES_BAD_KEY_MAT;

	if (keyMaterial == NULL)
		return AES_BAD_KEY_MAT;

	/* initialize key schedule: */
	for (i = 0; i < keyLen / 8; i++) {
		j = keyMaterial[i];
		k[i % 4][i / 4] = j;
	}

	rijndaelKeySched(k, keyLen, blockLen, W);
#ifdef MV_AES_DEBUG
	{
		MV_U8 *pW = &W[0][0][0];
		int x;

		mvOsPrintf("Expended Key: size = %d\n", sizeof(W));
		for (i = 0; i < sizeof(W); i++)
			mvOsPrintf("%02x ", pW[i]);

		for (i = 0; i < MAXROUNDS + 1; i++) {
			mvOsPrintf("\n Round #%02d: ", i);
			for (x = 0; x < MAXBC; x++)
				mvOsPrintf("%02x%02x%02x%02x ", W[i][0][x], W[i][1][x], W[i][2][x], W[i][3][x]);
			mvOsPrintf("\n");
		}
	}
#endif /* MV_AES_DEBUG */
	switch (keyLen) {
	case 128:
		rounds = 10;
		KC = 4;
		break;
	case 192:
		rounds = 12;
		KC = 6;
		break;
	case 256:
		rounds = 14;
		KC = 8;
		break;
	default:
		return (-1);
	}

	for (i = 0; i < MAXBC; i++)
		for (j = 0; j < 4; j++)
			expandedKey[i * 4 + j] = W[rounds][j][i];

	for (; i < KC; i++)
		for (j = 0; j < 4; j++)
			expandedKey[i * 4 + j] = W[rounds - 1][j][i + MAXBC - KC];

	return 0;
}

int aesBlockEncrypt128(MV_U8 mode, MV_U8 *IV, MV_U8 *expandedKey, int keyLen,
		       MV_U32 *plain, int numBlocks, MV_U32 *cipher)
{
	int i, j, t;
	MV_U8 block[4][MAXBC];
	int rounds;
	char *input, *outBuffer;

	input = (char *)plain;
	outBuffer = (char *)cipher;

	/* check parameter consistency: */
	if ((expandedKey == NULL) || ((keyLen != 128) && (keyLen != 192) && (keyLen != 256)))
		return AES_BAD_KEY_MAT;

	if ((mode != MODE_ECB && mode != MODE_CBC))
		return AES_BAD_CIPHER_STATE;

	switch (keyLen) {
	case 128:
		rounds = 10;
		break;
	case 192:
		rounds = 12;
		break;
	case 256:
		rounds = 14;
		break;
	default:
		return (-3);	/* this cannot happen */
	}

	switch (mode) {
	case MODE_ECB:
		for (i = 0; i < numBlocks; i++) {
			for (j = 0; j < 4; j++) {
				for (t = 0; t < 4; t++)
					/* parse input stream into rectangular array */
					block[t][j] = input[16 * i + 4 * j + t] & 0xFF;
			}
			rijndaelEncrypt128(block, (MV_U8(*)[4][MAXBC]) expandedKey, rounds);
			for (j = 0; j < 4; j++) {
				/* parse rectangular array into output ciphertext bytes */
				for (t = 0; t < 4; t++)
					outBuffer[16 * i + 4 * j + t] = (MV_U8) block[t][j];
			}
		}
		break;

	case MODE_CBC:
		for (j = 0; j < 4; j++) {
			for (t = 0; t < 4; t++)
				/* parse initial value into rectangular array */
				block[t][j] = IV[t + 4 * j] & 0xFF;
		}
		for (i = 0; i < numBlocks; i++) {
			for (j = 0; j < 4; j++) {
				for (t = 0; t < 4; t++)
					/* parse input stream into rectangular array and exor with
					   IV or the previous ciphertext */
					block[t][j] ^= input[16 * i + 4 * j + t] & 0xFF;
			}
			rijndaelEncrypt128(block, (MV_U8(*)[4][MAXBC]) expandedKey, rounds);
			for (j = 0; j < 4; j++) {
				/* parse rectangular array into output ciphertext bytes */
				for (t = 0; t < 4; t++)
					outBuffer[16 * i + 4 * j + t] = (MV_U8) block[t][j];
			}
		}
		break;

	default:
		return AES_BAD_CIPHER_STATE;
	}

	return 0;
}

int aesBlockDecrypt128(MV_U8 mode, MV_U8 *IV, MV_U8 *expandedKey, int keyLen,
		       MV_U32 *srcData, int numBlocks, MV_U32 *dstData)
{
	int i, j, t;
	MV_U8 block[4][MAXBC];
	MV_U8 iv[4][MAXBC];
	int rounds;
	char *input, *outBuffer;

	input = (char *)srcData;
	outBuffer = (char *)dstData;

	if (expandedKey == NULL)
		return AES_BAD_KEY_MAT;

	/* check parameter consistency: */
	if (keyLen != 128 && keyLen != 192 && keyLen != 256)
		return AES_BAD_KEY_MAT;

	if ((mode != MODE_ECB && mode != MODE_CBC))
		return AES_BAD_CIPHER_STATE;

	switch (keyLen) {
	case 128:
		rounds = 10;
		break;
	case 192:
		rounds = 12;
		break;
	case 256:
		rounds = 14;
		break;
	default:
		return (-3);	/* this cannot happen */
	}

	switch (mode) {
	case MODE_ECB:
		for (i = 0; i < numBlocks; i++) {
			for (j = 0; j < 4; j++) {
				for (t = 0; t < 4; t++) {
					/* parse input stream into rectangular array */
					block[t][j] = input[16 * i + 4 * j + t] & 0xFF;
				}
			}
			rijndaelDecrypt128(block, (MV_U8(*)[4][MAXBC]) expandedKey, rounds);
			for (j = 0; j < 4; j++) {
				/* parse rectangular array into output ciphertext bytes */
				for (t = 0; t < 4; t++)
					outBuffer[16 * i + 4 * j + t] = (MV_U8) block[t][j];
			}
		}
		break;

	case MODE_CBC:
		/* first block */
		for (j = 0; j < 4; j++) {
			for (t = 0; t < 4; t++) {
				/* parse input stream into rectangular array */
				block[t][j] = input[4 * j + t] & 0xFF;
				iv[t][j] = block[t][j];
			}
		}
		rijndaelDecrypt128(block, (MV_U8(*)[4][MAXBC]) expandedKey, rounds);

		for (j = 0; j < 4; j++) {
			/* exor the IV and parse rectangular array into output ciphertext bytes */
			for (t = 0; t < 4; t++) {
				outBuffer[4 * j + t] = (MV_U8) (block[t][j] ^ IV[t + 4 * j]);
				IV[t + 4 * j] = iv[t][j];
			}
		}

		/* next blocks */
		for (i = 1; i < numBlocks; i++) {
			for (j = 0; j < 4; j++) {
				for (t = 0; t < 4; t++) {
					/* parse input stream into rectangular array */
					iv[t][j] = input[16 * i + 4 * j + t] & 0xFF;
					block[t][j] = iv[t][j];
				}
			}
			rijndaelDecrypt128(block, (MV_U8(*)[4][MAXBC]) expandedKey, rounds);

			for (j = 0; j < 4; j++) {
				/* exor previous ciphertext block and parse rectangular array
				   into output ciphertext bytes */
				for (t = 0; t < 4; t++) {
					outBuffer[16 * i + 4 * j + t] = (MV_U8) (block[t][j] ^ IV[t + 4 * j]);
					IV[t + 4 * j] = iv[t][j];
				}
			}
		}
		break;

	default:
		return AES_BAD_CIPHER_STATE;
	}

	return 0;
}