/************************************************************************/
/*  Copyright (C), 2016-2020, [IT], 保留所有权利;
/*  模 块 名：AES算法加解密工具类;
/*  描    述：AES又称Rijndael，是高级加密标准的数据加密解密算法;
/*
/*  版    本：[V];	
/*  作    者：[IT];
/*  日    期：[6/3/2016];
/*
/*
/*  注    意：使用宏__AES_IMC_SPEEDUP__能加快生成轮密钥（AesCipher::makeRoundKey）,但意味着更多4K内存被使用。
/*			  默认情况下,禁用此加速,如果确实使用加速,开启宏定义即可;
/*
/*  修改记录：[IT];
/*  修改日期：;
/*  修改版本：;
/*  修改内容：;
/************************************************************************/
#ifndef __AES_CIPHER_H__
#define __AES_CIPHER_H__

#include <cstring>
#include <climits>
#include <stdexcept>

//#define __AES_IMC_SPEEDUP__	// 加密生成轮密钥的宏;

class AesCipher
{
public: 
	// AES的工作模式（同样是其他分组（块）加密算法的工作模式）
	// ECB：电子密码本模式(Electronic Codebook mode);
	// CBC：密码块链接模式(Cipher Block Chaining mode);
	// CFB：密码反馈模式(Cipher Feedback mode);
	// OFB：输出反馈模式(Output Feedback mode);
	// CTR：计数器模式(Counter mode);
	enum { ECB = 0, CBC = 1, CFB = 2, OFB = 3, CTR = 4 };
	enum { DEFAULT_KEY_LENGTH = 16 };      // 默认密钥长度16字节;
	enum { BLOCK_SIZE = 16 };              // 分组(块)大小;

public:
	AesCipher() : m_bKey(false), m_bIv(false), m_bCounter(false) { }

	// 将用户输入的密钥扩展成轮密钥（轮计算时用到的密钥）;
	void makeRoundKey(IN const char *key, IN int keyLength = DEFAULT_KEY_LENGTH);

	// 为链工作模式(CBC,CFB和OFB)设置初始向量(IV);
	void setIV(const char *iv) {
		if (iv != 0) {
			m_bIv = true;
			memcpy(m_iv, iv, BLOCK_SIZE);
		} else
			m_bIv = false;
	}

	// 为CTR工作模式设置初始计数器;
	// 参数：counter,为CTR(只适应128bits或16Bytes)而设置的初始计数器;
	// 注意：低地址字节对应计数器的低位;
	void setCounter(const char *counter) {
		if (counter != 0) {
			m_bCounter = true;
			memcpy(m_counter, counter, BLOCK_SIZE);
		} else
			m_bCounter = false;
	}

	// 加密多个明文块;
	void encrypt(IN const char *in, OUT char *result, IN size_t num, IN int mode = CBC);
	// 解密多个密文块;
	void decrypt(IN const char *in, OUT char *result, IN size_t num, IN int mode = CBC);
	// 用明文长度获取对应的密文长度(返回密文字节数);
	static int calculateCipherLen(size_t plaintextlen) {
		return (plaintextlen / BLOCK_SIZE + (plaintextlen % BLOCK_SIZE ? 1 : 0) + 1) * 16;
	}

	// 加密任意长度的明文;
	void encryptString(IN const char *in, OUT char *result, IN size_t num, IN int mode = CBC);

	// 解密密文;
	int decryptString(IN const char *in, OUT char *result, IN size_t num, IN int mode = CBC);
	
	// 获取密钥长度;
	int getKeyLength() {
		if (!m_bKey)
			return 0;
		return m_keyLength;
	}

	// 获取轮次数;
	int getRounds() {
		if (!m_bKey)
			return 0;
		return m_rounds;
	}

	// private types and functions
private:
	typedef unsigned __int32 u32_t;
	typedef unsigned char    u8_t;

	// 加密一个块大小的明文,这个块假定是AES的块大小（128bits）;
	void encryptBlock(IN const char *in, OUT char *result);

	// 解密一个块大小的密文,这个块假定是AES的块大小（128bits）;
	void decryptBlock(IN const char *in, OUT char *result);

	// 块的异或(XOR)运行;
	static void xorBlock(char *block1, const char *block2) {
		for (int i = 0; i < BLOCK_SIZE; ++i)
			block1[i] ^= block2[i]; 
	}

	// 计数器增加（Increase）1;
	bool incrCounter(char *counter);

private:
	enum { MAX_ROUNDS = 14 };        // maximum number of round
	enum { MAX_KWC = 8 };            // key's maximum number of word
	enum { BWC = 4 };                // block's number of word (BLOCK_SIZE / 4)

	// Substitute Bytes Transformation Array, S-Box
	static const u8_t sm_sbox[256];
	// Inverse Substitute Bytes Transformation Array, Inverse S-Box
	static const u8_t sm_isbox[256];
	// Two-in-one Transform Array: substitute byte & MixColumns transformation
	static const u32_t  sm_te1[256];
	static const u32_t  sm_te2[256];
	static const u32_t  sm_te3[256];
	static const u32_t  sm_te4[256];
	// Two-in-one Transform Array: inverse substitute byte & inverse MixColumns transformation
	static const u32_t  sm_td1[256];
	static const u32_t  sm_td2[256];
	static const u32_t  sm_td3[256];
	static const u32_t  sm_td4[256];

#ifdef __AES_IMC_SPEEDUP__
	// Inverse MixColumns Transformation Array
	static const u32_t  sm_imc1[256];          
	static const u32_t  sm_imc2[256];
	static const u32_t  sm_imc3[256];
	static const u32_t  sm_imc4[256];
#else
	// Inverse MixColumns Transformation Function
	unsigned char mul(unsigned char a, unsigned char b);
	u32_t invMixColumn(u32_t w);   
#endif 

	// Round Constant Array
	static const u8_t sm_rcon[52];           

	bool m_bKey;                        // key initialization flag
	u32_t m_ke[MAX_ROUNDS + 1][BWC];    // encryption round key
	u32_t m_kd[MAX_ROUNDS + 1][BWC];    // decryption round key
	int m_keyLength;                    // length of key
	int m_rounds;                       // number of rounds
	// initial vector (IV) block, for chain mode
	bool m_bIv;                         // IV initialization flag
	char m_iv[BLOCK_SIZE];
	// initial counter for CTR mode.
	bool m_bCounter;                    // counter initialization flag
	char m_counter[BLOCK_SIZE];
};

#endif  //__AES_CIPHER_H__