//
// Created by UnknownObject on 2022/10/22.
// Copyright (c) 2022 UnknownNetworkService. All rights reserved.
//

//数据处理算法
//这个算法不做额外说明，实际比赛中算法会发生变化

#include "main_car_aes.h"

namespace uns
{
	BinaryByteArray uns::MC_AES::CreateTextArray(std::string str)
	{
		if (str.size() > 16)
			return nullptr;
		BinaryByteArray arr = new ByteArray[4];
		for (int i = 0; i < 4; i++)
			arr[i] = new byte[4];
		if (str.size() != 16)
		{
			int fix_bit = 16 - str.size();
			while (str.size() < 16)
				str.push_back(fix_bit);
		}
		for (int i = 0; i < str.size(); i += 4)
		{
			for (int j = i; j < i + 4; j++)
				arr[j - i][i / 4] = str[j];
		}
		return arr;
	}
	
	word uns::MC_AES::Byte2Word(std::string str_4)
	{
		word result = 0;
		result |= ((unsigned long long)str_4[0] << 24);
		result |= ((unsigned long long)str_4[1] << 16);
		result |= ((unsigned long long)str_4[2] << 8);
		result |= str_4[3];
		return result;
	}
	
	WordArray uns::MC_AES::CreateKeyArray(std::string key)
	{
		if (key.size() != 16)
			return nullptr;
		WordArray arr = new word[4];
		for (int i = 0; i < 16; i += 4)
			arr[i / 4] = Byte2Word(key.substr(i, 4));
		return arr;
	}
	
	word uns::MC_AES::CircleLeft(word rw)
	{
		word high = rw << 8;
		word low = rw >> 24;
		return high | low;
	}
	
	ByteArray uns::MC_AES::MoveLeft(ByteArray bytes, int bits)
	{
		if (bits == 0)
			return bytes;
		else
		{
			bits = bits % 4;
			while (bits--)
			{
				byte index = bytes[0];
				for (int i = 0; i < 3; i++)
					bytes[i] = bytes[i + 1];
				bytes[3] = index;
			}
			return bytes;
		}
	}
	
	word uns::MC_AES::SubWord(word sw)
	{
		word temp;
		for (int i = 0; i < 32; i += 8)
		{
			int row = sw[i + 7] * 8 + sw[i + 6] * 4 + sw[i + 5] * 2 + sw[i + 4];
			int col = sw[i + 3] * 8 + sw[i + 2] * 4 + sw[i + 1] * 2 + sw[i];
			byte val = S_Box[row][col];
			for (int j = 0; j < 8; ++j)
				temp[i + j] = val[j];
		}
		return temp;
	}
	
	WordArray uns::MC_AES::KeyExpansion(WordArray fix_key)
	{
		WordArray keys = new word[44];
		for (int i = 0; i < 44; i++)
		{
			if (i < 4)
				keys[i] = fix_key[i];
			else if ((i % 4) != 0)
				keys[i] = keys[i - 4] ^ keys[i - 1];
			else
				keys[i] = keys[i - 4] ^ SubWord(CircleLeft(keys[i - 1])) ^ Rcon[(i / 4) - 1];
		}
		delete[] fix_key;
		return keys;
	}
	
	BinaryByteArray uns::MC_AES::BytesReplace(BinaryByteArray bytes)
	{
		for (int i = 0; i < 4; i++)
			for (int j = 0; j < 4; j++)
			{
				int row = bytes[i][j][7] * 8 + bytes[i][j][6] * 4 + bytes[i][j][5] * 2 + bytes[i][j][4];
				int col = bytes[i][j][3] * 8 + bytes[i][j][2] * 4 + bytes[i][j][1] * 2 + bytes[i][j][0];
				bytes[i][j] = S_Box[row][col];
			}
		return bytes;
	}
	
	BinaryByteArray uns::MC_AES::ByteShift(BinaryByteArray bytes)
	{
		for (int i = 0; i < 4; i++)
			bytes[i] = MoveLeft(bytes[i], i);
		return bytes;
	}
	
	BinaryByteArray uns::MC_AES::ColumnMix(BinaryByteArray bytes)
	{
		for(int i = 0; i < 4; i++)
			for (int j = 0; j < 4; j++)
				bytes[i][j] = bytes[i][j].to_ullong() + Columns[i][j].to_ulong();
		return bytes;
	}
	
	byte uns::MC_AES::SplitWord(word w, int index)
	{
		if ((index < 0) || (index >= 4))
			return 0;
		WordSpliter spliter;
		spliter._word = w.to_ullong();
		return spliter._result[index];
	}
	
	WordArray uns::MC_AES::GetCurrentKeys(WordArray keys, int times)
	{
		if ((times < 0) || (times >= 10))
			return nullptr;
		WordArray result = new word[4];
		for (int i = 4; i <= 7; i++)
			result[i - 4] = keys[i + times];
		return result;
	}
	
	BinaryByteArray uns::MC_AES::ColumnMix(BinaryByteArray bytes, WordArray keys)
	{
		for (int i = 0; i < 4; i++)
		{
			for (int j = 0; j < 4; j++)
				bytes[j][i] ^= SplitWord(keys[i], 3 - j).to_ullong();
		}
		delete[] keys;
		return bytes;
	}
	
	ByteArray uns::MC_AES::GetAlarmCode(BinaryByteArray bytes)
	{
		WordArray wcode = new word[4];
		ByteArray bcode = new byte[6];
		memset(wcode, 0, sizeof(wcode));
		memset(bcode, 0, sizeof(bcode));
		for (int i = 0; i < 4; i++)
			for (int j = 0; j < 4; j++)
				wcode[i] = wcode[i].to_ullong() + bytes[j][i].to_ullong();
		bcode[0] = S_Box[0x0][0x0];
		bcode[1] = S_Box[0xF][0xF];
		for (int i = 0; i < 4; i++)
			bcode[i + 2] = SplitWord(wcode[i], 0);
		delete[] wcode;
		return bcode;
	}
	
	ByteArray MC_AES::RunAES(std::string text, std::string key)
	{
		BinaryByteArray arr = CreateTextArray(text);
		WordArray warr = KeyExpansion(CreateKeyArray(key));
		for (int i = 0; i < 10; i++)
			arr = ColumnMix(ColumnMix(ByteShift(BytesReplace(arr))), GetCurrentKeys(warr, i));
		ByteArray code = GetAlarmCode(arr);
		delete[] arr;
		delete[] warr;
		return code;
	}
};

extern "C" JNIEXPORT
jstring JNICALL Java_com_uns_maincar_cpp_1interface_EnvTest_MainCarAESTest(JNIEnv *env, jclass _this)
{
	std::string version = "MainCar AES Found, Version: " + std::string(MAIN_CAR_AES_VERSION);
	return env->NewStringUTF(version.c_str());
}

extern "C" JNIEXPORT
jboolean JNICALL Java_com_uns_maincar_cpp_1interface_MainCarAES_RunAES(JNIEnv *env, jclass _this, jstring text, jbyteArray result)
{
	if(result == nullptr)
		return false;
	uns::MC_AES aes;
	uns::ByteArray arr = aes.RunAES(env->GetStringUTFChars(text, 0));
	jbyte barr[6] = { 0 };
	for(int i = 0; i < 6; i++)
		barr[i] = (jbyte)arr[i].to_ulong();
	env->SetByteArrayRegion(result, 0, 6, barr);
	delete[] arr;
	return true;
}