package sm1

import "encoding/binary"

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

const mask = 0xffffffff

func rol(x, i uint32) uint32 {
	return (x << i) | (x >> (32 - i))
}

func ror(x, i uint32) uint32 {
	return (x >> i) | (x << (32 - i))
}

// x,y int[4]
func r(y, x []uint32) {
	y[0] = rol(x[0], 1)
	y[1] = rol(x[1], 9)
	y[2] = rol(x[2], 17)
	y[3] = rol(x[3], 25)
}

// x,y int[4]
func rr(y, x []uint32) {
	y[0] = ror(x[0], 1)
	y[1] = ror(x[1], 9)
	y[2] = ror(x[2], 17)
	y[3] = ror(x[3], 25)
}

// x,y,rk int[4]
func j(y, x, rk []uint32) {
	y[0] = x[0] ^ rk[0]
	y[1] = x[1] ^ rk[1]
	y[2] = x[2] ^ rk[2]
	y[3] = x[3] ^ rk[3]
}

// x,y are [4]byte
func l(y []uint32, x []uint32) {
	y[0] = x[1] ^ x[2] ^ x[3]
	y[1] = x[1] ^ x[2] ^ x[0]
	y[2] = x[1] ^ x[0] ^ x[3]
	y[3] = x[0] ^ x[2] ^ x[3]
}

func s(y []uint32, x []uint32) {
	xBuf := make([]byte, 4)
	for i, u := range x {
		binary.LittleEndian.PutUint32(xBuf, u)
		for i := range xBuf {
			xBuf[i] = sbox[xBuf[i]]
		}
		y[i] = binary.LittleEndian.Uint32(xBuf)
	}
}

//enc = 1 for enc and 0 for dec
func round(out, m []uint32, rk []uint32, r func([]uint32, []uint32)) {
	y := make([]uint32, 4)
	j(out, m, rk)
	l(y, out)
	s(y, y) // L(S(y)) can do in one sbox change?
	l(out, y)
	r(out, out)
}

// input 16bytes to 4 uint32 in bigendian
func bytes2ints(out []uint32, in []byte) {
	out[0] = binary.BigEndian.Uint32(in[:4])
	out[1] = binary.BigEndian.Uint32(in[4:8])
	out[2] = binary.BigEndian.Uint32(in[8:12])
	out[3] = binary.BigEndian.Uint32(in[12:])
}

func ints2bytes(out []byte, in []uint32) {
	binary.BigEndian.PutUint32(out, in[0])
	binary.BigEndian.PutUint32(out[4:], in[1])
	binary.BigEndian.PutUint32(out[8:], in[2])
	binary.BigEndian.PutUint32(out[12:], in[3])
}

func block(out []byte, in []byte, key roundKey, r func([]uint32, []uint32)) {
	x := make([]uint32, 4)
	bytes2ints(x, in)
	r(x, x)
	for i := 0; i < numOfRound; i++ {
		round(x, x, key[4*i:], r)
	}

	j(x, x, key[4*numOfRound:])
	ints2bytes(out, x)
}

func encryptBlock(dst, src []byte, key []uint32) {
	block(dst, src, key, r)
}
func decryptBlock(dst, src []byte, key []uint32) {
	block(dst, src, key, rr)
}