#include <stdio.h>
//#include <stdint.h>
#include <x86intrin.h>
#include "sm4_ref.h"

//#define SM4_ROUND 32

#define ROTL32(x, n) (((x) << (n)) | ((x) >> (32 - (n))))
#define rotl32(t,num) (_mm_or_si128(_mm_slli_epi32(t,num),_mm_srli_epi32(t,32-num)))

static const uint32_t sm4_ck[32] = {
    0x00070E15, 0x1C232A31, 0x383F464D, 0x545B6269,
    0x70777E85, 0x8C939AA1, 0xA8AFB6BD, 0xC4CBD2D9,
    0xE0E7EEF5, 0xFC030A11, 0x181F262D, 0x343B4249,
    0x50575E65, 0x6C737A81, 0x888F969D, 0xA4ABB2B9,
    0xC0C7CED5, 0xDCE3EAF1, 0xF8FF060D, 0x141B2229,
    0x30373E45, 0x4C535A61, 0x686F767D, 0x848B9299,
    0xA0A7AEB5, 0xBCC3CAD1, 0xD8DFE6ED, 0xF4FB0209,
    0x10171E25, 0x2C333A41, 0x484F565D, 0x646B7279
};

static const uint8_t sm4_sbox[256] = {
    0xD6, 0x90, 0xE9, 0xFE, 0xCC, 0xE1, 0x3D, 0xB7,
    0x16, 0xB6, 0x14, 0xC2, 0x28, 0xFB, 0x2C, 0x05,
    0x2B, 0x67, 0x9A, 0x76, 0x2A, 0xBE, 0x04, 0xC3,
    0xAA, 0x44, 0x13, 0x26, 0x49, 0x86, 0x06, 0x99,
    0x9C, 0x42, 0x50, 0xF4, 0x91, 0xEF, 0x98, 0x7A,
    0x33, 0x54, 0x0B, 0x43, 0xED, 0xCF, 0xAC, 0x62,
    0xE4, 0xB3, 0x1C, 0xA9, 0xC9, 0x08, 0xE8, 0x95,
    0x80, 0xDF, 0x94, 0xFA, 0x75, 0x8F, 0x3F, 0xA6,
    0x47, 0x07, 0xA7, 0xFC, 0xF3, 0x73, 0x17, 0xBA,
    0x83, 0x59, 0x3C, 0x19, 0xE6, 0x85, 0x4F, 0xA8,
    0x68, 0x6B, 0x81, 0xB2, 0x71, 0x64, 0xDA, 0x8B,
    0xF8, 0xEB, 0x0F, 0x4B, 0x70, 0x56, 0x9D, 0x35,
    0x1E, 0x24, 0x0E, 0x5E, 0x63, 0x58, 0xD1, 0xA2,
    0x25, 0x22, 0x7C, 0x3B, 0x01, 0x21, 0x78, 0x87,
    0xD4, 0x00, 0x46, 0x57, 0x9F, 0xD3, 0x27, 0x52,
    0x4C, 0x36, 0x02, 0xE7, 0xA0, 0xC4, 0xC8, 0x9E,
    0xEA, 0xBF, 0x8A, 0xD2, 0x40, 0xC7, 0x38, 0xB5,
    0xA3, 0xF7, 0xF2, 0xCE, 0xF9, 0x61, 0x15, 0xA1,
    0xE0, 0xAE, 0x5D, 0xA4, 0x9B, 0x34, 0x1A, 0x55,
    0xAD, 0x93, 0x32, 0x30, 0xF5, 0x8C, 0xB1, 0xE3,
    0x1D, 0xF6, 0xE2, 0x2E, 0x82, 0x66, 0xCA, 0x60,
    0xC0, 0x29, 0x23, 0xAB, 0x0D, 0x53, 0x4E, 0x6F,
    0xD5, 0xDB, 0x37, 0x45, 0xDE, 0xFD, 0x8E, 0x2F,
    0x03, 0xFF, 0x6A, 0x72, 0x6D, 0x6C, 0x5B, 0x51,
    0x8D, 0x1B, 0xAF, 0x92, 0xBB, 0xDD, 0xBC, 0x7F,
    0x11, 0xD9, 0x5C, 0x41, 0x1F, 0x10, 0x5A, 0xD8,
    0x0A, 0xC1, 0x31, 0x88, 0xA5, 0xCD, 0x7B, 0xBD,
    0x2D, 0x74, 0xD0, 0x12, 0xB8, 0xE5, 0xB4, 0xB0,
    0x89, 0x69, 0x97, 0x4A, 0x0C, 0x96, 0x77, 0x7E,
    0x65, 0xB9, 0xF1, 0x09, 0xC5, 0x6E, 0xC6, 0x84,
    0x18, 0xF0, 0x7D, 0xEC, 0x3A, 0xDC, 0x4D, 0x20,
    0x79, 0xEE, 0x5F, 0x3E, 0xD7, 0xCB, 0x39, 0x48
};

static const uint32_t sm4_fk[4] = {
  0xA3B1BAC6, 0x56AA3350, 0x677D9197, 0xB27022DC
};

static uint32_t load_uint32_t_be(const uint8_t *b, uint32_t n)
{
	return ((uint32_t)b[4 * n] << 24) |
		((uint32_t)b[4 * n + 1] << 16) |
		((uint32_t)b[4 * n + 2] << 8) |
		((uint32_t)b[4 * n + 3]);
}

static void store_uint32_t_be(uint32_t v, uint8_t *b)
{
	b[0] = (uint8_t)(v >> 24);
	b[1] = (uint8_t)(v >> 16);
	b[2] = (uint8_t)(v >> 8);
	b[3] = (uint8_t)(v);
}

void sm4_key_schedule(const uint8_t key[], uint32_t rk[])
{
	uint32_t t, x, k[36];
	int i;

	for (i = 0; i < 4; i++)
	{
		k[i] = load_uint32_t_be(key, i) ^ sm4_fk[i];
	}

	/* T' */
	for (i = 0; i < SM4_ROUND; ++i)
	{
		x = k[i + 1] ^ k[i + 2] ^ k[i + 3] ^ sm4_ck[i];

		/* Nonlinear operation tau */
		t = ((uint32_t)sm4_sbox[(uint8_t)(x >> 24)]) << 24 |
			((uint32_t)sm4_sbox[(uint8_t)(x >> 16)]) << 16 |
			((uint32_t)sm4_sbox[(uint8_t)(x >> 8)]) << 8 |
			((uint32_t)sm4_sbox[(uint8_t)(x)]);

		/* Linear operation L' */
		k[i + 4] = k[i] ^ (t ^ ROTL32(t, 13) ^ ROTL32(t, 23));
		rk[i] = k[i + 4];
	}
}

static uint32_t sm4_s(uint32_t x)
{
	return (((uint32_t)sm4_sbox[x >> 24]) << 24) | (((uint32_t)sm4_sbox[(x >> 16) & 0x0ff]) << 16) | (((uint32_t)sm4_sbox[(x >> 8) & 0x0ff]) << 8) | ((uint32_t)sm4_sbox[x & 0x0ff]);
	
}

static __m128i sm4_t(__m128i x)
{
	uint32_t t[4];
	_mm_storeu_si128((__m128i *)t, x);
	for (int i = 0; i < 4; i++)
	{
		t[i] = sm4_s(t[i]);
	}
	x = _mm_loadu_si128((__m128i *)t);
	return _mm_xor_si128(_mm_xor_si128(_mm_xor_si128(_mm_xor_si128(x, rotl32(x, 2)), rotl32(x, 10)), rotl32(x, 18)), rotl32(x, 24));
}

void sm4_encrypt(const uint32_t rk[SM4_ROUND], const uint8_t *plaintext, uint8_t *ciphertext)
{
	const __m128i flp = _mm_set_epi32(0x0C0D0E0F, 0x08090A0B, 0x04050607, 0x00010203);
	//const __m128i flp = _mm_set_epi32(0x00010203, 0x04050607, 0x08090A0B, 0x0C0D0E0F);
	//const __m128i flp = _mm_set_epi16(0x0001, 0x0203, 0x0405, 0x0607, 0x0809, 0x0a0b, 0x0c0d, 0x0e0f);
	
	uint32_t *p32;
	uint32_t k, v[4];
	p32 = (uint32_t *)plaintext;
	__m128i t0, t1, t2, t3, x;
	t0 = _mm_set_epi32(p32[12], p32[8], p32[4], p32[0]);
	t0 = _mm_shuffle_epi8(t0, flp);
	t1 = _mm_set_epi32(p32[13], p32[9], p32[5], p32[1]);
	t1 = _mm_shuffle_epi8(t1, flp);
	t2 = _mm_set_epi32(p32[14], p32[10], p32[6], p32[2]);
	t2 = _mm_shuffle_epi8(t2, flp);
	t3 = _mm_set_epi32(p32[15], p32[11], p32[7], p32[3]);
	t3 = _mm_shuffle_epi8(t3, flp);
	/*
	uint32_t tt[4];
	_mm_storeu_si128((__m128i *)tt, t0);
	printf("%08x %08x %08x %08x", tt[0],tt[1],tt[2],tt[3]);
	printf("\n");
	*/
	for (int i = 0; i < SM4_ROUND; ++i)
	{
		k = rk[i];
		//x = t1 ^ t2 ^ t3 ^ _mm_set_epi32(k, k, k, k);
		x = _mm_xor_si128(t1, _mm_xor_si128(t2, _mm_xor_si128(t3, _mm_set1_epi32(k))));
		t0 = _mm_xor_si128(t0, sm4_t(x));
		x = t0;
		t0 = t1;
		t1 = t2;
		t2 = t3;
		t3 = x;
		/*
		uint32_t tt[4];
		_mm_storeu_si128((__m128i *)tt, x);
		printf("%d %x\n", i, tt[0]);
		*/
	}

	p32 = (uint32_t *)ciphertext;

	_mm_storeu_si128((__m128i *) v, _mm_shuffle_epi8(t3, flp));
	p32[0] = v[0];
	p32[4] = v[1];
	p32[8] = v[2];
	p32[12] = v[3];

	_mm_storeu_si128((__m128i *) v, _mm_shuffle_epi8(t2, flp));
	p32[1] = v[0];
	p32[5] = v[1];
	p32[9] = v[2];
	p32[13] = v[3];

	_mm_storeu_si128((__m128i *) v, _mm_shuffle_epi8(t1, flp));
	p32[2] = v[0];
	p32[6] = v[1];
	p32[10] = v[2];
	p32[14] = v[3];

	_mm_storeu_si128((__m128i *) v, _mm_shuffle_epi8(t0, flp));
	p32[3] = v[0];
	p32[7] = v[1];
	p32[11] = v[2];
	p32[15] = v[3];
}
/*
int main() 
{
	const uint8_t key[16] = {   0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef,0xfe,0xdc,0xba,0x98,0x76,0x54,0x32,0x10 };
	const uint8_t input[64] = { 0x67,0x45,0x23,0x01,0xef,0xcd,0xab,0x89,0x98,0xba,0xdc,0xfe,0x10,0x32,0x54,0x76,
								0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef,0xfe,0xdc,0xba,0x98,0x76,0x54,0x32,0x10 ,
								0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef,0xfe,0xdc,0xba,0x98,0x76,0x54,0x32,0x10 ,
								0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef,0xfe,0xdc,0xba,0x98,0x76,0x54,0x32,0x10 };
	uint8_t output[64];
	uint32_t rk[32];
	sm4_key_schedule(key, rk);
	sm4_encrypt(rk, input, output);
	for (int i = 0; i < 64; i++) 
	{
		printf("%02x ", output[i]);
		if (i % 16 == 15)
			printf("\n");
	}
	getchar();
	return 0;

}
*/
