#include "aes128.h"

// S盒 (Substitution Box)
static const u8 sbox[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};

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

// 轮常量
static const u8 rcon[10] = {0x01, 0x02, 0x04, 0x08, 0x10,
							0x20, 0x40, 0x80, 0x1b, 0x36};

// 伽罗瓦域乘法 (xtime) - 确保正确
static u8 xtime(u8 x) { return (x << 1) ^ (((x >> 7) & 1) * 0x1b); }

// 关键修复：AES状态矩阵是列优先排列！
// 内存布局应该是：s00, s10, s20, s30, s01, s11, s21, s31, ...
// 而不是行优先

// 轮密钥加 - 无需修改
void add_round_key(u8 *state, const u8 *round_key) {
	for (u8 i = 0; i < 16; i++) {
		state[i] ^= round_key[i];
	}
}

// 字节替换 - 无需修改
void sub_bytes(u8 *state) {
	for (u8 i = 0; i < 16; i++) {
		state[i] = sbox[state[i]];
	}
}

// 行移位 - 修复为列优先排列
void shift_rows(u8 *state) {
	u8 temp;

	// AES状态矩阵（列优先）：
	// 列0: state[0], state[1], state[2], state[3]   -> 行0,1,2,3
	// 列1: state[4], state[5], state[6], state[7]   -> 行0,1,2,3
	// 列2: state[8], state[9], state[10], state[11] -> 行0,1,2,3
	// 列3: state[12], state[13], state[14], state[15] -> 行0,1,2,3

	// 第1行（索引0,4,8,12）不移位
	// 第2行（索引1,5,9,13）左移1位
	temp = state[1];
	state[1] = state[5];
	state[5] = state[9];
	state[9] = state[13];
	state[13] = temp;

	// 第3行（索引2,6,10,14）左移2位
	temp = state[2];
	state[2] = state[10];
	state[10] = temp;
	temp = state[6];
	state[6] = state[14];
	state[14] = temp;

	// 第4行（索引3,7,11,15）左移3位
	temp = state[3];
	state[3] = state[15];
	state[15] = state[11];
	state[11] = state[7];
	state[7] = temp;
}

// 列混合 - 修复为正确的实现
void mix_columns(u8 *state) {
	u8 a[4], b[4];

	// 处理每一列
	for (u8 col = 0; col < 4; col++) {
		// 获取当前列的4个字节
		a[0] = state[col * 4 + 0];
		a[1] = state[col * 4 + 1];
		a[2] = state[col * 4 + 2];
		a[3] = state[col * 4 + 3];

		// 计算混合后的值
		b[0] = xtime(a[0]) ^ xtime(a[1]) ^ a[1] ^ a[2] ^ a[3];
		b[1] = a[0] ^ xtime(a[1]) ^ xtime(a[2]) ^ a[2] ^ a[3];
		b[2] = a[0] ^ a[1] ^ xtime(a[2]) ^ xtime(a[3]) ^ a[3];
		b[3] = xtime(a[0]) ^ a[0] ^ a[1] ^ a[2] ^ xtime(a[3]);

		// 存回结果
		state[col * 4 + 0] = b[0];
		state[col * 4 + 1] = b[1];
		state[col * 4 + 2] = b[2];
		state[col * 4 + 3] = b[3];
	}
}

// 伽罗瓦域乘法 (用于逆列混合)
static u8 gf_mul(u8 a, u8 b) {
	u8 p = 0;
	u8 high_bit;

	while (b) {
		if (b & 1) {
			p ^= a;
		}
		high_bit = a & 0x80;
		a <<= 1;
		if (high_bit) {
			a ^= 0x1b; // x^8 + x^4 + x^3 + x + 1
		}
		b >>= 1;
	}
	return p;
}

// 逆字节替换
void inv_sub_bytes(u8 *state) {
	for (u8 i = 0; i < 16; i++) {
		state[i] = inv_sbox[state[i]];
	}
}

// 逆行移位 - 与shift_rows对应
void inv_shift_rows(u8 *state) {
	u8 temp;

	// 第2行右移1位
	temp = state[13];
	state[13] = state[9];
	state[9] = state[5];
	state[5] = state[1];
	state[1] = temp;

	// 第3行右移2位
	temp = state[2];
	state[2] = state[10];
	state[10] = temp;
	temp = state[6];
	state[6] = state[14];
	state[14] = temp;

	// 第4行右移3位
	temp = state[7];
	state[7] = state[11];
	state[11] = state[15];
	state[15] = state[3];
	state[3] = temp;
}

// 逆列混合 - 正确实现
void inv_mix_columns(u8 *state) {
	u8 a[4], b[4];

	// 处理每一列
	for (u8 col = 0; col < 4; col++) {
		// 获取当前列的4个字节
		a[0] = state[col * 4 + 0];
		a[1] = state[col * 4 + 1];
		a[2] = state[col * 4 + 2];
		a[3] = state[col * 4 + 3];

		// 计算逆混合后的值
		b[0] = gf_mul(0x0e, a[0]) ^ gf_mul(0x0b, a[1]) ^ gf_mul(0x0d, a[2]) ^
			   gf_mul(0x09, a[3]);
		b[1] = gf_mul(0x09, a[0]) ^ gf_mul(0x0e, a[1]) ^ gf_mul(0x0b, a[2]) ^
			   gf_mul(0x0d, a[3]);
		b[2] = gf_mul(0x0d, a[0]) ^ gf_mul(0x09, a[1]) ^ gf_mul(0x0e, a[2]) ^
			   gf_mul(0x0b, a[3]);
		b[3] = gf_mul(0x0b, a[0]) ^ gf_mul(0x0d, a[1]) ^ gf_mul(0x09, a[2]) ^
			   gf_mul(0x0e, a[3]);

		// 存回结果
		state[col * 4 + 0] = b[0];
		state[col * 4 + 1] = b[1];
		state[col * 4 + 2] = b[2];
		state[col * 4 + 3] = b[3];
	}
}

// 密钥扩展 - 修复为标准实现
static void key_expansion(aes128_context *ctx) {
	u8 temp[4];
	u8 i, rcon_iter = 0;

	// 复制原始密钥到第一轮密钥
	memcpy(ctx->round_keys[0], ctx->key, 16);

	// 生成10轮密钥
	for (i = 1; i <= 10; i++) {
		// 复制前一轮最后一列
		temp[0] = ctx->round_keys[i - 1][12];
		temp[1] = ctx->round_keys[i - 1][13];
		temp[2] = ctx->round_keys[i - 1][14];
		temp[3] = ctx->round_keys[i - 1][15];

		// 字循环 (RotWord)
		u8 t = temp[0];
		temp[0] = temp[1];
		temp[1] = temp[2];
		temp[2] = temp[3];
		temp[3] = t;

		// 字节替换 (SubWord)
		temp[0] = sbox[temp[0]];
		temp[1] = sbox[temp[1]];
		temp[2] = sbox[temp[2]];
		temp[3] = sbox[temp[3]];

		// 轮常量异或
		temp[0] ^= rcon[i - 1];

		// 生成新的第一列
		ctx->round_keys[i][0] = ctx->round_keys[i - 1][0] ^ temp[0];
		ctx->round_keys[i][1] = ctx->round_keys[i - 1][1] ^ temp[1];
		ctx->round_keys[i][2] = ctx->round_keys[i - 1][2] ^ temp[2];
		ctx->round_keys[i][3] = ctx->round_keys[i - 1][3] ^ temp[3];

		// 生成其他列
		for (u8 j = 1; j < 4; j++) {
			ctx->round_keys[i][j * 4 + 0] =
				ctx->round_keys[i][(j - 1) * 4 + 0] ^
				ctx->round_keys[i - 1][j * 4 + 0];
			ctx->round_keys[i][j * 4 + 1] =
				ctx->round_keys[i][(j - 1) * 4 + 1] ^
				ctx->round_keys[i - 1][j * 4 + 1];
			ctx->round_keys[i][j * 4 + 2] =
				ctx->round_keys[i][(j - 1) * 4 + 2] ^
				ctx->round_keys[i - 1][j * 4 + 2];
			ctx->round_keys[i][j * 4 + 3] =
				ctx->round_keys[i][(j - 1) * 4 + 3] ^
				ctx->round_keys[i - 1][j * 4 + 3];
		}
	}
}

// 设置密钥
void aes128_set_key(aes128_context *ctx, const u8 *key) {
	memcpy(ctx->key, key, AES128_KEY_LEN);
	key_expansion(ctx);
}

// 加密单个块
void aes128_encrypt_block(aes128_context *ctx, const u8 *input, u8 *output) {
	// 复制输入到状态
	memcpy(ctx->state, input, AES_BLOCK_SIZE);

	// 初始轮密钥加
	add_round_key(ctx->state, ctx->round_keys[0]);

	// 9轮标准轮
	for (u8 round = 1; round <= 9; round++) {
		sub_bytes(ctx->state);
		shift_rows(ctx->state);
		mix_columns(ctx->state);
		add_round_key(ctx->state, ctx->round_keys[round]);
	}

	// 最终轮 (没有mix_columns)
	sub_bytes(ctx->state);
	shift_rows(ctx->state);
	add_round_key(ctx->state, ctx->round_keys[10]);

	// 复制结果到输出
	memcpy(output, ctx->state, AES_BLOCK_SIZE);
}

// ECB模式加密48字节数据
void aes128_ecb_encrypt_48bytes(aes128_context *ctx, const u8 *input,
								u8 *output) {
	// 分3块加密 (48 = 16 * 3)
	aes128_encrypt_block(ctx, input, output);
	aes128_encrypt_block(ctx, input + 16, output + 16);
	aes128_encrypt_block(ctx, input + 32, output + 32);
}

// 解密单个块
void aes128_decrypt_block(aes128_context *ctx, const u8 *input, u8 *output) {
	// 复制输入到状态
	memcpy(ctx->state, input, AES_BLOCK_SIZE);

	// 初始轮密钥加（使用最后一轮密钥）
	add_round_key(ctx->state, ctx->round_keys[10]);

	// 9轮标准解密轮
	for (u8 round = 9; round >= 1; round--) {
		inv_shift_rows(ctx->state);
		inv_sub_bytes(ctx->state);
		add_round_key(ctx->state, ctx->round_keys[round]);
		inv_mix_columns(ctx->state);
	}

	// 最终轮（没有inv_mix_columns）
	inv_shift_rows(ctx->state);
	inv_sub_bytes(ctx->state);
	add_round_key(ctx->state, ctx->round_keys[0]);

	// 复制结果到输出
	memcpy(output, ctx->state, AES_BLOCK_SIZE);
}

// ECB模式解密48字节数据
void aes128_ecb_decrypt_48bytes(aes128_context *ctx, const u8 *input,
								u8 *output) {
	// 分3块解密 (48 = 16 * 3)
	aes128_decrypt_block(ctx, input, output);
	aes128_decrypt_block(ctx, input + 16, output + 16);
	aes128_decrypt_block(ctx, input + 32, output + 32);
}

// 验证函数：加密后解密，对比原始数据
u8 aes128_verify_encryption(aes128_context *ctx, const u8 *original_data,
							size_t length) {
	u8 temp_ciphertext[48];
	u8 temp_plaintext[48];
	u8 result = 1; // 1表示验证成功，0表示失败

	// 1. 加密原始数据
	aes128_ecb_encrypt_48bytes(ctx, original_data, temp_ciphertext);

	// 2. 解密密文
	aes128_ecb_decrypt_48bytes(ctx, temp_ciphertext, temp_plaintext);

	// 3. 对比原始数据和解密后的数据
	for (size_t i = 0; i < length; i++) {
		if (original_data[i] != temp_plaintext[i]) {
			result = 0; // 验证失败
			break;
		}
	}

	return result;
}