#include <stdint.h>

typedef unsigned char   uint8_t;
typedef unsigned long   uint32_t;

/////////////////////////////////////////////////
// SM4 加密/解密算法
//  X4 = X0 ⊕ T(X1 ⊕ X2 ⊕ X3 ⊕ rk)
//      ⊕ 是异或操作
//      rk 是秘钥
//      T 是查表 + 离散操作
//      X0 ~ X3 是原始输入
//
/////////////////////////////////////////////////  
//  秘钥 rk 的生成
//  1、原始秘钥： MK = 0123456789abcdeffedcba9876543210
//  2、拆成4个32位：
//      MK0 = 0x01234567
//      MK1 = 0x89abcdef
//      MK2 = 0xfedcba98
//      MK3 = 0x76543210
//  3、SM4固定4个常量：
//      FK0 = 0xa3b1bac6
//      FK1 = 0x56aa3350
//      FK2 = 0x677d9197
//      FK3 = 0xb27022dc
//  4、计算
//      K0 = MK0 ⊕ FK0 = 0xa292ffa1
//      K1 = MK1 ⊕ FK1 = 0xdf01febf
//      K2 = MK2 ⊕ FK2 = 0x99a12b0f
//      K3 = MK3 ⊕ FK3 = 0xc42410cc
//  5、CK 常量（32个）
//      CK0 = 0x00070e15
//      CK1 = 0x1c232a31
//      CK2 = 0x383f464d
//      ...
//  6、32 轮密钥生成公式
//      K[i+4] = K[i] ⊕ T'(K[i+1] ⊕ K[i+2] ⊕ K[i+3] ⊕ CK[i])
//      rk[i]  = K[i+4]
//  7、T'
//      T'(x) = S盒 + L'
//      S盒是将32bit的tmp值拆分成4个8bit的值, 然后每个 8bit通过数组查表映射成另外一个值, 
//          最后再将这4个映射后的值按原位置拼接起来
//      L'(x) = x ⊕ (x <<< 13) ⊕ (x <<< 23)
//          （<<< 是循环左移）
//  8、举例算出 rk0
//      tmp = K1 ⊕ K2 ⊕ K3 ⊕ CK0
//          = 0xdf01febf ⊕ 0x99a12b0f ⊕ 0xc42410cc ⊕ 0x00070e15
//          = 0x8283cb69
//      过 S盒 82 83 cb 69 ==> 0x2abe5c82  (映射数组可以对照下面代码的 Sbox[256])
//      L' 变换 ： L'(x) = x ⊕ (x<<<13) ⊕ (x<<<23) = 0xa00746fb
//             其中 x      = 0x2abe5c82
//                  x<<<13 = 0xcb904557
//                  x<<<23 = 0x41295f2e
//      计算rk0 : rk0 = K0 ⊕ T'(tmp)
//                    = 0xa292ffa1 ⊕ 0xa00746fb
//                    = 0xf12186f9  
//
/////////////////////////////////////////////////
// 第一轮加密过程: X4 = X0 ⊕ T(X1 ⊕ X2 ⊕ X3 ⊕ rk)
//  初值:
//      X0 = 0x01234567
//      X1 = 0x89abcdef
//      X2 = 0xfedcba98
//      X3 = 0x76543210
//      rk0 = 0xf12186f9
//  Step1：先 XOR
//      tmp = X1 ⊕ X2 ⊕ X3 ⊕ rk0
//          = 0x89abcdef ⊕ 0xfedcba98 ⊕ 0x76543210 ⊕ 0xf12186f9
//          = 0xf000c39e
//  Step2：过 S盒
//      f0 00 c3 9e ==> 0x18d65ab8
//  Step3：线性变换 L
//      L(x) = x ⊕ (x<<<2) ⊕ (x<<<10) ⊕ (x<<<18) ⊕ (x<<<24) = 0xf01d6538
//           其中 : x      = 0x18d65ab8
//                  x<<<2  = 0x63596ae0
//                  x<<<10 = 0x596ae063
//                  x<<<18 = 0x6ae06359
//                  x<<<24 = 0xb818d65a
//  Step4：得到 X4
//      X4 = X0 ⊕ T(tmp)
//         = 0x01234567 ⊕ 0xf01d6538
//         = 0x27fad345  
//
///////////////////////////////////////////////// 

// ================= 高性能 SM4 实现 =================
// S盒
static const uint8_t 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
};

// 循环左移 (x <<< n)
#define ROTL(x,n)   (((x) << (n)) | ((x) >> (32-(n))))

// ================= 优化点1：全部 inline =================
// 32bit 数据 A 过 S 盒
static inline uint32_t tau(uint32_t A) {
    return (Sbox[A>>24]<<24) 
        | (Sbox[(A>>16)&0xff]<<16) 
        | (Sbox[(A>>8)&0xff]<<8) 
        | (Sbox[A&0xff]);
}

// 数据的线性变化 L(x) = x ^ (x<<<2) ^ (x<<<10) ^ (x<<<18) ^ (x<<<24)
static inline uint32_t L(uint32_t B) {
    return B ^ ROTL(B,2) ^ ROTL(B,10) ^ ROTL(B,18) ^ ROTL(B,24);
}

// 秘钥的线性变化 L'(x) = x ^ (x<<<13) ^ (x<<<23)
static inline uint32_t Lp(uint32_t B) {
    return B ^ ROTL(B,13) ^ ROTL(B,23);
}

static inline uint32_t T(uint32_t x)    { return L(tau(x));  }
static inline uint32_t Tp(uint32_t x)   { return Lp(tau(x)); }

// ================= 常量 =================
static const uint32_t FK[4] = {
    0xa3b1bac6,0x56aa3350,0x677d9197,0xb27022dc
};

static const uint32_t 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
};

// ================= 优化点2：密钥扩展（无log） =================
//  MK[4] : 原始秘钥 
//  rk[32] : 输出 32轮的秘钥
void sm4_key_schedule(uint32_t MK[4], uint32_t rk[32]) {
    uint32_t K[4];

    // 计算 K0 ~ K3
    K[0] = MK[0] ^ FK[0];
    K[1] = MK[1] ^ FK[1];
    K[2] = MK[2] ^ FK[2];
    K[3] = MK[3] ^ FK[3];

    // 循环 32 次计算 : K[i+4] = K[i] ^ T'(K[i+1] ^ K[i+2] ^ K[i+3] ^ CK[i])
    for (int i = 0; i < 32; i++) {
        // K[i+1] ^ K[i+2] ^ K[i+3] ^ CK[i]
        uint32_t tmp = K[1] ^ K[2] ^ K[3] ^ CK[i];
        // T'()
        uint32_t t = Tp(tmp);
        // K[i] ^ T'()
        uint32_t newK = K[0] ^ t;
        // 保存本轮秘钥
        rk[i] = newK;

        // 滚动（避免36数组）
        K[0] = K[1];
        K[1] = K[2];
        K[2] = K[3];
        K[3] = newK;
    }
}

// ================= 优化点3：加密（寄存器滚动） =================
// X[4] : 原始数据 
// rk[32] : 32轮计算每轮秘钥
// out[4] : 加密结果
void sm4_encrypt(uint32_t X[4], uint32_t rk[32], uint32_t out[4]) {
    uint32_t x0=X[0], x1=X[1], x2=X[2], x3=X[3];

    // 32轮计算 X4 = X0 ^ T(X1 ^ X2 ^ X3 ^ rk)
    for(int i=0; i<32; i++) {
        uint32_t tmp = x1 ^ x2 ^ x3 ^ rk[i];
        uint32_t t = T(tmp);
        uint32_t newX = x0 ^ t;

        // 滚动寄存器
        x0=x1; x1=x2; x2=x3; x3=newX;
    }

    // 输出反序
    out[0]=x3; out[1]=x2; out[2]=x1; out[3]=x0;
}

// ================= 优化点4：解密 =================
// X[4]   : 加密数据 
// rk[32] : 32轮计算每轮秘钥
// out[4] : 解密结果
void sm4_decrypt(uint32_t C[4], uint32_t rk[32], uint32_t out[4]) {
    uint32_t x0=C[0], x1=C[1], x2=C[2], x3=C[3];

    // 32轮计算 X4 = X0 ⊕ T(X1 ⊕ X2 ⊕ X3 ⊕ rk')
    for(int i=0;i<32;i++){
        uint32_t tmp = x1 ^ x2 ^ x3 ^ rk[31-i];
        uint32_t t = T(tmp);
        uint32_t newX = x0 ^ t;

        // 滚动寄存器
        x0=x1; x1=x2; x2=x3; x3=newX;
    }

    // 输出反序
    out[0]=x3; out[1]=x2; out[2]=x1; out[3]=x0;
}

// ====== 测试 ======
int main(){
    uint32_t MK[4] = {0x01234567,0x89abcdef,0xfedcba98,0x76543210};
    uint32_t P[4]  = {0x01234567,0x89abcdef,0xfedcba98,0x76543210};

    uint32_t rk[32],C[4],D[4];

    // MK 秘钥扩展为 32轮的秘钥 rk[32]
    sm4_key_schedule(MK,rk);
    // 编码 -> C[]
    sm4_encrypt(P,rk,C);
    // 解码 -> D[]
    sm4_decrypt(C,rk,D);

    // 打印结果
    printf("Plain : %08x %08x %08x %08x\n",P[0],P[1],P[2],P[3]);
    printf("Cipher: %08x %08x %08x %08x\n",C[0],C[1],C[2],C[3]);
    printf("Dec   : %08x %08x %08x %08x\n",D[0],D[1],D[2],D[3]);

    return 0;
}