#include "Crypto.h"


#define MD5_A 0x67452301
#define MD5_B 0xefcdab89
#define MD5_C 0x98badcfe
#define MD5_D 0x10325476

const char MD5_S16[] = "0123456789abcdef";

const unsigned int T[] = {
	0xd76aa478,0xe8c7b756,0x242070db,0xc1bdceee,
	0xf57c0faf,0x4787c62a,0xa8304613,0xfd469501,
	0x698098d8,0x8b44f7af,0xffff5bb1,0x895cd7be,
	0x6b901122,0xfd987193,0xa679438e,0x49b40821,
	0xf61e2562,0xc040b340,0x265e5a51,0xe9b6c7aa,
	0xd62f105d,0x02441453,0xd8a1e681,0xe7d3fbc8,
	0x21e1cde6,0xc33707d6,0xf4d50d87,0x455a14ed,
	0xa9e3e905,0xfcefa3f8,0x676f02d9,0x8d2a4c8a,
	0xfffa3942,0x8771f681,0x6d9d6122,0xfde5380c,
	0xa4beea44,0x4bdecfa9,0xf6bb4b60,0xbebfbc70,
	0x289b7ec6,0xeaa127fa,0xd4ef3085,0x04881d05,
	0xd9d4d039,0xe6db99e5,0x1fa27cf8,0xc4ac5665,
	0xf4292244,0x432aff97,0xab9423a7,0xfc93a039,
	0x655b59c3,0x8f0ccc92,0xffeff47d,0x85845dd1,
	0x6fa87e4f,0xfe2ce6e0,0xa3014314,0x4e0811a1,
	0xf7537e82,0xbd3af235,0x2ad7d2bb,0xeb86d391 };

const unsigned int MD5_OFFSET[] = {
	7,12,17,22,7,12,17,22,7,12,17,22,7,12,17,22,
	5,9,14,20,5,9,14,20,5,9,14,20,5,9,14,20,
	4,11,16,23,4,11,16,23,4,11,16,23,4,11,16,23,
	6,10,15,21,6,10,15,21,6,10,15,21,6,10,15,21
};

// 填充字符串
vector<unsigned int> padding(string src, unsigned int& len)
{
	// 以512位,64个字节为一组
	unsigned int num = ((src.length() + 8) / 64) + 1;
	vector<unsigned int> rec(num * 16);
	len = num * 16;
	for (unsigned int i = 0; i < src.length(); i++) {
		// 一个unsigned int对应4个字节，保存4个字符信息
		rec[i >> 2] |= (int)(src[i]) << ((i % 4) * 8);
	}
	// 补充1000...000
	rec[src.length() >> 2] |= (0x80 << ((src.length() % 4) * 8));
	// 填充原文长度
	rec[rec.size() - 2] = (src.length() << 3);
	return rec;
}
// 整理输出
string md5_format(unsigned int num)
{
	string res = "";
	unsigned int base = 1 << 8;
	for (int i = 0; i < 4; i++) {
		string tmp = "";
		unsigned int b = (num >> (i * 8)) % base & 0xff;
		for (int j = 0; j < 2; j++) {
			tmp = MD5_S16[b % 16] + tmp;
			b /= 16;
		}
		res += tmp;
	}
	return res;
}

// F函数
unsigned int F(unsigned int b, unsigned int c, unsigned int d)
{
	return (b & c) | ((~b) & d);
}
// G函数
unsigned int G(unsigned int b, unsigned int c, unsigned int d)
{
	return (b & d) | (c & (~d));
}
// H函数
unsigned int H(unsigned int b, unsigned int c, unsigned int d)
{
	return b ^ c ^ d;
}
// I函数
unsigned int I(unsigned int b, unsigned int c, unsigned int d)
{
	return c ^ (b | (~d));
}
// 移位操作函数
unsigned int shift(unsigned int a, unsigned int n)
{
	return (a << n) | (a >> (32 - n));
}

// 循环压缩
void md5_encode(unsigned int* X, int size, unsigned int& tempA, unsigned int& tempB, unsigned int& tempC, unsigned int& tempD)
{
	unsigned int a = tempA,
		b = tempB,
		c = tempC,
		d = tempD,
		rec = 0,
		g, k;
	for (int i = 0; i < 64; i++) {
		if (i < 16) {
			// F迭代
			g = F(b, c, d);
			k = i;
		} else if (i < 32) {
			// G迭代
			g = G(b, c, d);
			k = (1 + 5 * i) % 16;
		} else if (i < 48) {
			// H迭代
			g = H(b, c, d);
			k = (5 + 3 * i) % 16;
		} else {
			// I迭代
			g = I(b, c, d);
			k = (7 * i) % 16;
		}
		rec = d;
		d = c;
		c = b;
		b = b + shift(a + g + X[k] + T[i], MD5_OFFSET[i]);
		a = rec;
	}
	tempA += a;
	tempB += b;
	tempC += c;
	tempD += d;
}

string Crypto::md5(string src)
{
	unsigned int tempA = MD5_A;
	unsigned int tempB = MD5_B;
	unsigned int tempC = MD5_C;
	unsigned int tempD = MD5_D;
	unsigned int len = 0;
	vector<unsigned int> rec = padding(src, len);
	for (unsigned int i = 0; i < len / 16; i++) {
		unsigned int num[16];
		for (int j = 0; j < 16; j++) {
			num[j] = rec[i * 16 + j];
		}
		md5_encode(num, 16, tempA, tempB, tempC, tempD);
	}
	return md5_format(tempA) + md5_format(tempB) + md5_format(tempC) + md5_format(tempD);
}

#include "common/Utils.h"

static std::vector<uint8_t> hexToByte(const std::string& hexstr)
{
	std::vector<uint8_t> v;
	v.reserve(hexstr.size() / 2);
	for (std::string::const_iterator it = hexstr.begin(); it != hexstr.end(); it += 2)
	{
		v.push_back(static_cast<char>(strtol(std::string(it, it + 2).c_str(), nullptr, 16)));
	}
	return v;
}

std::string Crypto::gmssl_sm2Encrypt(std::string pubKeyHex, std::string priKeyHex, std::string plainByte)
{
	std::vector<uint8_t> pub_key = hexToByte(pubKeyHex);

	std::vector<uint8_t> pri_keyBuf = hexToByte(priKeyHex);
	sm2_z256_t pri_key {0};
	sm2_z256_from_bytes(pri_key, pri_keyBuf.data());

	SM2_KEY sm2k;
	sm2_key_set_public_key(&sm2k, (SM2_Z256_POINT*)pub_key.data()); // +1 跳过 0x04 标记
	sm2_key_set_private_key(&sm2k, pri_key);
	sm2_key_print(stderr, 0, 4, "SM2_KEY", &sm2k);

	//std::vector<uint8_t> plainText(plainByte.begin(), plainByte.end());
	std::vector<uint8_t> ciphertext(SM2_MAX_CIPHERTEXT_SIZE);
	size_t ciphertext_len = 0;

	int ret = sm2_encrypt(&sm2k, (const uint8_t*)plainByte.data(), plainByte.size(), ciphertext.data(), &ciphertext_len);
	return std::string(ciphertext.begin(), ciphertext.begin()+ciphertext_len);
}


std::string Crypto::gmssl_sm2Decrypt(std::string pubKeyHex, std::string priKeyHex, std::string cipherByte)
{
	std::vector<uint8_t> pub_key = hexToByte(pubKeyHex);

	std::vector<uint8_t> pri_keyBuf = hexToByte(priKeyHex);
	sm2_z256_t pri_key {0};
	sm2_z256_from_bytes(pri_key, pri_keyBuf.data());

	SM2_KEY sm2k;
	sm2_key_set_public_key(&sm2k, (SM2_Z256_POINT*)pub_key.data()); // +1 跳过 0x04 标记
	sm2_key_set_private_key(&sm2k, pri_key);
	sm2_key_print(stderr, 0, 4, "SM2_KEY", &sm2k);

	//std::string ciphertHex = "44d37f6c644f22508b5387663e3a174384917eba7ce3b2a5bb9f79e191d40dfc6c0cfaf1819af562c1ee1bffcc55d2fb1da8114527e3a6993aaadda12a26b4b6f688ecd5d7700c8bd82a069ab4b30e0b499eef82d1dfeea261d2fd19d1f0ae95a4b3fd";
	//Utils::toHex(ciphertext, ciphertext_len);
	//std::vector<uint8_t> buf = hexToByte(ciphertHex);
	std::vector<uint8_t> plainText(SM2_MAX_PLAINTEXT_SIZE);
	size_t len = 0;
	int ret = sm2_decrypt(&sm2k, (const uint8_t*)cipherByte.data(), cipherByte.size(), plainText.data(), &len);
	return std::string(plainText.begin(), plainText.begin() + len);
}


string Crypto::sm3(std::string data)
{
	SM3_CTX ctx;
	sm3_init(&ctx);
	sm3_update(&ctx, (const unsigned char*)data.data(), data.length());
	std::vector<unsigned char> digest(SM3_DIGEST_SIZE);
	sm3_finish(&ctx, digest.data());
	std::stringstream ss;
	for (int i = 0; i < digest.size(); ++i)
	{
		ss << std::hex << std::setw(2) << std::setfill('0') << (int)digest[i];
	}
	return ss.str();
}

std::vector<uint8_t> Crypto::sm4EcbEncrypt(const std::vector<uint8_t>& key, const std::vector<uint8_t>& plaintext)
{
	SM4_KEY sm4_key;
	std::vector<uint8_t> ciphertext(plaintext.size()); // ECB模式密文长度与明文相同（需填充时另说）

	sm4_set_encrypt_key(&sm4_key, key.data());
	// 假设明文已是分组整数倍或已处理填充
	for (size_t i = 0; i < plaintext.size(); i += 16) {
		sm4_encrypt(&sm4_key, &plaintext[i], &ciphertext[i]);
	}
	return ciphertext;
}

std::vector<uint8_t> Crypto::sm4EcbDecrypt(const std::vector<uint8_t>& key, const std::vector<uint8_t>& ciphertext) {
	SM4_KEY sm4_key;
	std::vector<uint8_t> plaintext(ciphertext.size());

	sm4_set_decrypt_key(&sm4_key, key.data());
	for (size_t i = 0; i < ciphertext.size(); i += 16) {
		sm4_encrypt(&sm4_key, &ciphertext[i], &plaintext[i]); // 注意解密时也是用sm4_encrypt函数，但密钥是解密密钥
	}
	return plaintext;
}


// 错误处理函数
static void HandleOpensslError() 
{
	ERR_print_errors_fp(stderr);
	unsigned long err = ERR_get_error();
}

// SM2 加密函数
std::string Crypto::ssl_sm2Encrypt(const std::string& public_key_hex, const std::string& plaintext)
{
	EVP_PKEY* pubkey = Crypto::create_sm2_public_key_from_hex(public_key_hex);

	EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new(pubkey, NULL);
	if (!ctx) HandleOpensslError();

	if (EVP_PKEY_encrypt_init(ctx) <= 0) HandleOpensslError();

	// 获取加密后数据的长度
	size_t ciphertext_len;
	if (EVP_PKEY_encrypt(ctx, NULL, &ciphertext_len, (const uint8_t*)plaintext.data(), plaintext.size()) <= 0) {
		HandleOpensslError();
	}

	// 执行加密
	std::vector<unsigned char> ciphertext(ciphertext_len);
	if (EVP_PKEY_encrypt(ctx, ciphertext.data(), &ciphertext_len, (const uint8_t*)plaintext.data(), plaintext.size()) <= 0) {
		HandleOpensslError();
	}

	// 清理资源
	EVP_PKEY_CTX_free(ctx);
	EVP_PKEY_free(pubkey);
	// 清理 OpenSSL
	EVP_cleanup();
	ERR_free_strings();

	return std::string(ciphertext.begin(), ciphertext.begin() + ciphertext_len);
}

// SM2 解密函数
std::string Crypto::ssl_sm2Decrypt(const std::string& private_key_hex, const std::string& ciphertext)
{
	EVP_PKEY* privkey = Crypto::create_sm2_private_key_from_hex(private_key_hex);
	EVP_PKEY_CTX* ctx = EVP_PKEY_CTX_new(privkey, NULL);
	if (!ctx) HandleOpensslError();

	if (EVP_PKEY_decrypt_init(ctx) <= 0) HandleOpensslError();

	// 获取解密后数据的长度
	size_t plaintext_len;
	if (EVP_PKEY_decrypt(ctx, NULL, &plaintext_len, (const uint8_t*)ciphertext.data(), ciphertext.size()) <= 0) {
		HandleOpensslError();
	}

	// 执行解密
	std::vector<unsigned char> plaintext(plaintext_len);
	if (EVP_PKEY_decrypt(ctx, plaintext.data(), &plaintext_len, (const uint8_t*)ciphertext.data(), ciphertext.size()) <= 0) {
		HandleOpensslError();
	}

	EVP_PKEY_CTX_free(ctx);
	EVP_PKEY_free(privkey);

	// 清理 OpenSSL
	EVP_cleanup();
	ERR_free_strings();

	return std::string(plaintext.begin(), plaintext.begin() + plaintext_len);
}

// 将16进制字符串转换为字节数组
static std::vector<unsigned char> hexToBytes(const std::string& hex) {
	std::vector<unsigned char> bytes;
	for (size_t i = 0; i < hex.length(); i += 2) 
	{
		std::string byteString = hex.substr(i, 2);
		unsigned char byte = static_cast<unsigned char>(strtol(byteString.c_str(), NULL, 16));
		bytes.push_back(byte);
	}

	return bytes;
}

// 从16进制字符串创建SM2私钥
EVP_PKEY* Crypto::create_sm2_private_key_from_hex(const std::string& private_key_hex)
{
	// 将16进制私钥转换为字节
	std::vector<unsigned char> private_key_bytes = hexToBytes(private_key_hex);

	// 创建BIGNUM
	BIGNUM* privBN = BN_bin2bn(private_key_bytes.data(), private_key_bytes.size(), NULL);
	if (!privBN) HandleOpensslError();

	// 创建EC_KEY并设置曲线和私钥
	EC_KEY* ecKey = EC_KEY_new_by_curve_name(NID_sm2);
	if (!ecKey) HandleOpensslError();

	if (!EC_KEY_set_private_key(ecKey, privBN)) {
		HandleOpensslError();
	}

	// 计算公钥
	const EC_GROUP* group = EC_KEY_get0_group(ecKey);
	EC_POINT* pub_point = EC_POINT_new(group);
	if (!pub_point) HandleOpensslError();

	if (!EC_POINT_mul(group, pub_point, privBN, NULL, NULL, NULL)) {
		HandleOpensslError();
	}

	if (!EC_KEY_set_public_key(ecKey, pub_point)) {
		HandleOpensslError();
	}

	// 创建EVP_PKEY
	EVP_PKEY* pkey = EVP_PKEY_new();
	if (!pkey) HandleOpensslError();

	if (!EVP_PKEY_set1_EC_KEY(pkey, ecKey)) {
		HandleOpensslError();
	}

	// 清理资源
	BN_free(privBN);
	EC_POINT_free(pub_point);
	EC_KEY_free(ecKey);

	return pkey;
}



// 从16进制字符串创建SM2公钥
EVP_PKEY* Crypto::create_sm2_public_key_from_hex(const std::string& public_key_hex) 
{
	// 将16进制公钥转换为字节
	std::vector<unsigned char> public_key_bytes = hexToBytes(public_key_hex);

	// 创建EC_KEY
	EC_KEY* ec_key = EC_KEY_new_by_curve_name(NID_sm2);
	if (!ec_key) HandleOpensslError();

	// 从字节创建公钥点
	const EC_GROUP* group = EC_KEY_get0_group(ec_key);
	EC_POINT* pub_point = EC_POINT_new(group);
	if (!pub_point) HandleOpensslError();

	if (!EC_POINT_oct2point(group, pub_point, public_key_bytes.data(), public_key_bytes.size(), NULL)) {
		HandleOpensslError();
	}

	if (!EC_KEY_set_public_key(ec_key, pub_point)) {
		HandleOpensslError();
	}

	// 创建EVP_PKEY
	EVP_PKEY* pkey = EVP_PKEY_new();
	if (!pkey) HandleOpensslError();

	if (!EVP_PKEY_set1_EC_KEY(pkey, ec_key)) {
		HandleOpensslError();
	}

	// 清理资源
	EC_POINT_free(pub_point);
	EC_KEY_free(ec_key);

	return pkey;
}