package sm4

import (
	"bytes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestSm4RoundKey(t *testing.T) {
	key, _ := hex.DecodeString("0123456789abcdeffedcba9876543210")
	enckey := make([]uint32, 32)
	deckey := make([]uint32, 32)
	expandKeyGo(key, enckey, deckey)
	for i := 0; i < 32; i++ {
		fmt.Printf("0x%02x, ", (enckey[i]>>24)&0xff)
		fmt.Printf("0x%02x, ", (enckey[i]>>16)&0xff)
		fmt.Printf("0x%02x, ", (enckey[i]>>8)&0xff)
		fmt.Printf("0x%02x, ", enckey[i]&0xff)
	}

}

func TestSM4StdBlock(t *testing.T) {
	key, _ := hex.DecodeString("0123456789abcdeffedcba9876543210")
	msg, _ := hex.DecodeString("0123456789abcdeffedcba9876543210")
	msg, _ = ECBEncrypt(msg, key, msg)
	enckey := make([]uint32, 32)
	deckey := make([]uint32, 32)
	expandKeyGo(key, enckey, deckey)
	for i := 0; i < 32; i++ {
		fmt.Printf("0x%02x, ", enckey[i])
	}

	waned, _ := hex.DecodeString("681edf34d206965e86b3e94f536e4246")
	assert.Equal(t, msg, waned)
}

func TestSm4Block(t *testing.T) {
	key, _ := hex.DecodeString("0123456789abcdeffedcba9876543210")

	// test 4 blocks for aesni
	msg, _ := hex.DecodeString("0123456789abcdeffedcba98765432100123456789abcdeffedcba98765432100123456789abcdeffedcba98765432100123456789abcdeffedcba9876543210")
	times := 1000000
	for i := 0; i < times; i++ {
		msg, _ = ECBEncrypt(msg, key, msg)
	}

	if hex.EncodeToString(msg) != "595298c7c6fd271f0402f804c33d3f66595298c7c6fd271f0402f804c33d3f66595298c7c6fd271f0402f804c33d3f66595298c7c6fd271f0402f804c33d3f66" {
		t.Log("SM4 Encrypt test Failed!")
		t.Log(hex.EncodeToString(msg))
		t.Fail()
	}

	for i := 0; i < times; i++ {
		msg, _ = ECBDecrypt(msg, key, msg)
	}

	if hex.EncodeToString(msg) != "0123456789abcdeffedcba98765432100123456789abcdeffedcba98765432100123456789abcdeffedcba98765432100123456789abcdeffedcba9876543210" {
		t.Log("SM4 Decrypt test Failed!")
		t.Log(hex.EncodeToString(msg))
		t.Fail()
	}
}

func BenchmarkSM4(b *testing.B) {
	key, _ := hex.DecodeString("0123456789abcdeffedcba9876543210")
	msg := make([]byte, 1024*1024)
	rand.Read(msg)
	b.ReportAllocs()
	b.SetBytes(int64(len(msg)))
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		msg, _ = ECBEncrypt(msg, key, msg)
	}

}

func TestECBSpeed(t *testing.T) {
	key := []byte("0123456789abcdef")
	data := make([]byte, 10*1024*1024) // 10 MB data
	rand.Read(data)

	start := time.Now()
	times := 100
	for i := 0; i < times; i++ {
		if _, err := EncryptECB(data, key, data); err != nil {
			panic(err)
		}
	}
	end := time.Now()
	elapsed := end.Sub(start)
	t.Log("SM4 Encrypt: ", int(float64(times*len(data))/float64(1024*1024)/elapsed.Seconds()), "MBps")

	start = time.Now()
	for i := 0; i < times; i++ {
		if _, err := DecryptECB(data, key, data); err != nil {
			panic(err)
		}
	}
	end = time.Now()
	elapsed = end.Sub(start)
	t.Log("SM4 Decrypt: ", int(float64(times*len(data))/float64(1024*1024)/elapsed.Seconds()), "MBps")
}

func TestReportAeadSpeed(t *testing.T) {
	key := []byte("0123456789abcdef")
	data := make([]byte, 10*1024)
	ad := make([]byte, 10*1024)
	rand.Read(data)
	rand.Read(ad)

	b, _ := NewCipher(key)
	aead, _ := cipher.NewGCM(b)

	dst := make([]byte, len(data)+aead.NonceSize())
	nonce := make([]byte, 12)

	start := time.Now()
	times := 10000
	for i := 0; i < times; i++ {
		dst = aead.Seal(dst[:0], nonce, data, ad)
	}
	end := time.Now()
	elapsed := end.Sub(start)
	t.Log("SM4 AEAD Seal: ", int(float64(times*len(data))/float64(1024*1024)/elapsed.Seconds()), "MBps")

	start = time.Now()
	for i := 0; i < times; i++ {
		if _, err := aead.Open(data[:0], nonce, dst, ad); err != nil {
			panic(err)
		}
	}
	end = time.Now()
	elapsed = end.Sub(start)
	t.Log("SM4 AEAD Open: ", int(float64(times*len(data))/float64(1024*1024)/elapsed.Seconds()), "MBps")
}

/////////////////////////////////////////////
//cbc test

type CBCKeyTest struct {
	key     []byte
	iv      []byte
	data    []byte
	encData []byte
}

var keyTests = []CBCKeyTest{
	{
		[]byte{
			0xE5, 0x26, 0xCD, 0x8D, 0xFC, 0x4E, 0x3B, 0xCA, 0xB0, 0xBA, 0x15, 0xF3, 0xC8, 0x22, 0xB6, 0xAE,
		},
		[]byte{
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		},
		[]byte{
			0xBB, 0x5D, 0xCE, 0x7C, 0x7E, 0x45, 0x3B, 0xFA, 0xC8, 0x4D, 0x67, 0x75, 0x81, 0xEE, 0xEC, 0x16,
			0xED, 0xA7, 0xE1, 0xAD, 0x4F, 0x94, 0x82, 0x82, 0x80, 0xDF, 0x1F, 0xC7, 0xA8, 0x9B, 0xB2, 0xBE,
		},
		[]byte{
			0x63, 0x9f, 0x4c, 0x4e, 0x52, 0x4a, 0x8c, 0x90, 0xfb, 0xd6, 0xc6, 0xf4, 0x2f, 0x18, 0xf3, 0xb5,
			0x28, 0x48, 0xae, 0xc0, 0xfc, 0xc0, 0x1b, 0x03, 0xff, 0x04, 0xc5, 0x50, 0x3c, 0x9a, 0xe6, 0x11,
		},
	},
	{
		[]byte{
			0x5F, 0x02, 0xD7, 0xE5, 0x47, 0x74, 0xE4, 0xE9, 0x43, 0x25, 0x68, 0xA3, 0x2F, 0x7E, 0x26, 0xAB,
		},
		[]byte{
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		},
		[]byte{
			0x3D, 0xAC, 0x76, 0x3E, 0x46, 0xB6, 0xF0, 0x49, 0x77, 0x4F, 0xA2, 0x2F, 0x8D, 0x51, 0xB4, 0xF4,
			0xD0, 0x96, 0xA7, 0xB5, 0x45, 0x05, 0xF6, 0xB3, 0x8F, 0x94, 0x55, 0xB1, 0xA4, 0x2E, 0x5D, 0xD6,
		},
		[]byte{
			0x4e, 0x3a, 0x19, 0x0d, 0x97, 0x10, 0xee, 0x4e, 0xcf, 0x87, 0xfa, 0xa5, 0xc3, 0x3b, 0x4b, 0x24,
			0xf3, 0x95, 0x84, 0xbb, 0x6e, 0xb2, 0x3a, 0x74, 0x26, 0xe9, 0x87, 0x65, 0x08, 0x58, 0xc6, 0x8d,
		},
	},
	{
		[]byte{
			0x0D, 0x55, 0x89, 0x44, 0x29, 0xB7, 0xDC, 0xA3, 0x8B, 0xE0, 0x2E, 0x98, 0x29, 0x63, 0xA5, 0x75,
		},
		[]byte{
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		},
		[]byte{
			0xB6, 0x97, 0x93, 0x80, 0x63, 0x68, 0xD2, 0x36, 0xF6, 0xD3, 0xD3, 0x5D, 0xBF, 0x5D, 0xAD, 0xA3,
			0xBF, 0xC0, 0x9F, 0x53, 0xC3, 0xAF, 0xF8, 0x48, 0x7E, 0xA4, 0xD9, 0x56, 0x60, 0x2A, 0xBA, 0xE7,
		},
		[]byte{
			0xaa, 0xeb, 0xf0, 0xec, 0xdf, 0x98, 0x72, 0xef, 0xde, 0x14, 0xc2, 0xe9, 0x06, 0x0d, 0x65, 0x06,
			0xba, 0x5b, 0x22, 0x1a, 0x7a, 0x8b, 0xf0, 0xe8, 0x52, 0x43, 0x78, 0xcc, 0xed, 0x71, 0x70, 0xd5,
		},
	},
	{
		[]byte{
			0x72, 0x31, 0x6E, 0x03, 0x4A, 0xDB, 0xEF, 0x7C, 0xC9, 0xD1, 0x7A, 0xA1, 0x1F, 0x25, 0x99, 0x09,
		},
		[]byte{
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		},
		[]byte{
			0xE4, 0xB4, 0x69, 0x8B, 0xDE, 0x75, 0x7D, 0xFF, 0xD8, 0xC6, 0xE6, 0x7F, 0xD7, 0x67, 0xBE, 0x62,
			0x08, 0xEC, 0xB7, 0xFC, 0x92, 0xFB, 0x77, 0xD0, 0x79, 0xC2, 0x69, 0xC5, 0x7E, 0x0A, 0xB6, 0x2B,
		},
		[]byte{
			0x90, 0x2c, 0x7c, 0x36, 0xfc, 0x6f, 0xc1, 0x3d, 0x09, 0x7d, 0xa5, 0xa8, 0x9b, 0x74, 0xc7, 0x05,
			0xf6, 0x3b, 0xc1, 0x15, 0x4b, 0x69, 0xad, 0xdb, 0x41, 0x8a, 0xf9, 0x94, 0xde, 0xed, 0x34, 0x3a,
		},
	},
	{
		[]byte{
			0xAC, 0xE8, 0x7E, 0x4D, 0x00, 0xD0, 0x55, 0xB5, 0x24, 0xB8, 0x5D, 0xAA, 0x23, 0xEE, 0x29, 0x18,
		},
		[]byte{
			0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
		},
		[]byte{
			0x03, 0x7B, 0x89, 0x8B, 0x3F, 0x53, 0x55, 0x57, 0x64, 0x3B, 0x88, 0x3E, 0x63, 0xF4, 0x17, 0x77,
			0xBA, 0x4D, 0xE7, 0xEF, 0xDE, 0xCB, 0x83, 0x0A, 0xC3, 0x77, 0x2A, 0x3D, 0x3B, 0x64, 0xBE, 0x92,
		},
		[]byte{
			0x36, 0x18, 0x2e, 0xc4, 0xe2, 0x73, 0x14, 0xd7, 0xb7, 0xb2, 0xfa, 0x5c, 0x0b, 0xa6, 0x83, 0xa3,
			0x2d, 0xb4, 0xe9, 0x81, 0x09, 0x5c, 0xee, 0xd5, 0xe0, 0x3c, 0xdd, 0xdf, 0x49, 0x28, 0x0f, 0x23,
		},
	},
}

func TestSm4Cbc(t *testing.T) {
	for _, test := range keyTests {
		output := make([]byte, len(test.data))
		iv := make([]byte, 16)
		copy(iv, test.iv)
		output, _ = CBCEncrypt(output, iv, test.key, test.data)
		// t.Log(test.iv)
		// t.Log(iv)
		if hex.EncodeToString(output) != hex.EncodeToString(test.encData) {
			t.Log("SM4 CBC Test Failed")
			t.Fail()
		}
		copy(iv, test.iv)
		output, _ = CBCDecrypt(output, iv, test.key, output)
		if hex.EncodeToString(output) != hex.EncodeToString(test.data) {
			t.Log("SM4 CBC Test Failed")
			t.Fail()
		}
	}
}

////////////////////////////////////////////
// ofb

func TestSm4OFB(t *testing.T) {
	var key = []byte{
		0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef,
		0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32, 0x10,
	}
	iv := make([]byte, len(key))
	var msg = []byte{
		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,
		0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef,
		0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32, 0x10,
	}

	encrypted := make([]byte, len(msg))
	decrypted := make([]byte, len(msg))
	encrypted, _ = OFBEncrypt(encrypted, iv, key, msg)

	// if hex.EncodeToString(encrypted) != "681edf34d206965e86b3e94f536e4246681edf34d206965e86b3e94f536e4246681edf34d206965e86b3e94f536e4246681edf34d206965e86b3e94f536e4246" {
	// t.Log(hex.EncodeToString(encrypted))
	// 	t.Fail()
	// }

	decrypted, _ = OFBDecrypt(decrypted, iv, key, encrypted)

	// if hex.EncodeToString(decrypted) != "0123456789abcdeffedcba98765432100123456789abcdeffedcba98765432100123456789abcdeffedcba98765432100123456789abcdeffedcba9876543210" {
	// t.Log(hex.EncodeToString(decrypted))
	// 	t.Fail()
	// }

	if hex.EncodeToString(decrypted) != hex.EncodeToString(msg) {
		t.Log("SM4 Decrypt test Failed!")
		t.Log(hex.EncodeToString(msg))
		t.Fail()
	}

}

func TestSm4Mac(t *testing.T) {
	key, _ := hex.DecodeString("55E26781FF0CB07963363B0B1880BD9F")
	input, _ := hex.DecodeString("4808F443E2F6BC0A7A0F0FB4A4B5375C3B7BCD235134DA5664F030296B81A60FB3A4DE1661EB3C3728416002F8888AC2")
	// input := make([]byte, 1000*16)
	iv, _ := hex.DecodeString("D43A478CEFAA5923E3288281D8DC7964")
	var sm4Mac []byte
	var err error
	cnt := 50000
	start := time.Now()
	for i := 0; i < cnt; i++ {
		sm4Mac, err = Mac(key, iv, input)
		if err != nil {
			t.Fatal("err:" + err.Error())
		}
	}
	end := time.Now()
	elapsed := end.Sub(start)
	fmt.Printf("Sign %d, used time: %d ms, %d pcs/s\n", cnt, elapsed.Milliseconds(), int(float64(cnt)/float64(elapsed.Milliseconds())*1000))

	fmt.Println(hex.EncodeToString(sm4Mac))
}

// sm4_GCM模式
var sm4GCMTests = []struct {
	key, nonce, plaintext, ad, result string
}{
	{
		"11754cd72aec309bf52f7687212e8957",
		"3c819d9a9bed087615030b65", // nonce should be 12 bytes.
		"plaintext",
		"additional message not need encrypt, empty is ok",
		"6111f78f2f82b913c20e333160bfec034c3720ac133a6203b1",
	},
}

func TestSM4GCM(t *testing.T) {
	for i, test := range sm4GCMTests {
		key, _ := hex.DecodeString(test.key)
		var sm4gcm cipher.AEAD
		var err error

		block, err := NewCipher(key)
		if err != nil {
			t.Fatal(err)
		}
		sm4gcm, err = cipher.NewGCM(block)

		if err != nil {
			t.Fatal(err)
		}

		nonce, _ := hex.DecodeString(test.nonce)
		plaintext := []byte(test.plaintext)

		ad := []byte(test.ad)
		fmt.Printf("plaintext:%x\n", plaintext)

		// Seal分配密文空间
		// ct := sm4gcm.Seal(nil, nonce, plaintext, ad)

		// 若有已分配好的密文空间buf
		// buf := make([]byte, len(plaintext)+sm4gcm.Overhead())
		// ct := sm4gcm.Seal(buf[:0], nonce, plaintext, ad)

		// 如果plaintext的cap足够，则可以重复利用
		for i := 0; i < sm4gcm.Overhead(); i++ {
			plaintext = append(plaintext, 0)
		}
		plaintext = plaintext[:len(plaintext)-sm4gcm.Overhead()]
		ct := sm4gcm.Seal(plaintext[:0], nonce, plaintext, ad)
		fmt.Printf("cipherText: %x\n", ct)

		if ctHex := hex.EncodeToString(ct); ctHex != test.result {
			t.Errorf("#%d: got %s, want %s", i, ctHex, test.result)
			continue
		}

		plaintext, err = sm4gcm.Open(ct[:0], nonce, ct, ad)
		if err != nil {
			t.Errorf("#%d: Open failed", i)
			continue
		}

		if !bytes.Equal(plaintext, []byte(test.plaintext)) {
			t.Errorf("#%d: plaintext's don't match: got %x vs %x", i, test.plaintext, plaintext)
			continue
		}

		// if ad, nonce, ct was changed, return err
		if len(ad) > 0 {
			ad[0] ^= 0x80
			if _, err := sm4gcm.Open(nil, nonce, ct, ad); err == nil {
				t.Errorf("#%d: Open was successful after altering additional data", i)
			}
			ad[0] ^= 0x80
		}

		nonce[0] ^= 0x80
		if _, err := sm4gcm.Open(nil, nonce, ct, ad); err == nil {
			t.Errorf("#%d: Open was successful after altering nonce", i)
		}
		nonce[0] ^= 0x80

		ct[0] ^= 0x80
		if _, err := sm4gcm.Open(nil, nonce, ct, ad); err == nil {
			t.Errorf("#%d: Open was successful after altering ciphertext", i)
		}
		ct[0] ^= 0x80
	}
}

func TestTTT(t *testing.T) {

	key := []byte{0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x98, 0x76, 0x54, 0x32, 0x10}
	iv := []byte{0xF0, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8, 0xF9, 0xFA, 0xFB, 0xFC, 0xFD, 0xFE, 0xFF}
	src := make([]byte, 2048)
	dst := make([]byte, 2048)
	for i := range src {
		src[i] = byte(i)
	}
	c, _ := NewCipher(key)
	ctr := cipher.NewCTR(c, iv)
	ctr.XORKeyStream(dst, src)
	for i, c := range dst {
		if i > 0 && i%16 == 0 {
			fmt.Println("")
		}
		fmt.Printf("0x%02x, ", c)

	}

}