package blockmode_test

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"testing"
	"time"

	"xdx.jelly/xgcl/grand"
	"xdx.jelly/xgcl/sm/sm4"
	"xdx.jelly/xgcl/utils/blockmode"
)

// 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 TestSM4Speed(t *testing.T) {
	key, err := hex.DecodeString("12345678123456781234567812345678")
	if err != nil {
		t.Fatal(err)
	}
	block, err := sm4.NewCipher(key)
	if err != nil {
		t.Fatal(err)
	}
	gcm, err := blockmode.NewGCM(blockmode.Wrap(block))
	if err != nil {
		t.Fatal(err)
	}
	nonce, err := hex.DecodeString("123456781234567812345678")
	if err != nil {
		t.Fatal(err)
	}

	plaintext := grand.GetRandom(1024 * 1024)
	ad := []byte("additional message not need encrypt, empty is ok")
	ct2 := make([]byte, 0, 1024*1024+gcm.Overhead())
	cnt := 100
	start := time.Now()
	for i := 0; i < cnt; i++ {
		err = gcm.EncryptInit(nonce)
		if err != nil {
			t.Fatal(err)
		}
		gcm.SpecifyADD(ad)
		ct2, err := gcm.EncryptUpdate(ct2, plaintext)
		if err != nil {
			t.Fatal(err)
		}

		ct2, err = gcm.EncryptFinal(ct2)
		if err != nil {
			t.Fatal(err)
		}
	}
	end := time.Now()
	elapsed := end.Sub(start)
	fmt.Printf("%f Bps\n", float64(len(plaintext)*cnt)/1024/1024*1000/float64(elapsed.Milliseconds()))

}

func TestSM4GCM(t *testing.T) {
	for i, test := range sm4GCMTests {
		key, _ := hex.DecodeString(test.key)

		block, _ := sm4.NewCipher(key)
		gcm, err := blockmode.NewGCM(blockmode.Wrap(block))
		if err != nil {
			t.Fatal(err)
		}

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

		ad := []byte(test.ad)

		// 提前分配好空间
		for i := 0; i < gcm.Overhead(); i++ {
			plaintext = append(plaintext, 0)
		}
		plaintext = plaintext[:len(plaintext)-gcm.Overhead()]

		ct := gcm.Seal(plaintext[:0], nonce, plaintext, ad)

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

		plaintext, err = gcm.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 := gcm.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 := gcm.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 := gcm.Open(nil, nonce, ct, ad); err == nil {
			t.Errorf("#%d: Open was successful after altering ciphertext", i)
		}
		ct[0] ^= 0x80
	}
}

func TestGcmUpdate(t *testing.T) {
	test := sm4GCMTests[0]
	key, _ := hex.DecodeString(test.key)

	block, _ := sm4.NewCipher(key)
	gcm, err := blockmode.NewGCM(blockmode.Wrap(block))
	if err != nil {
		t.Fatal(err)
	}

	nonce, _ := hex.DecodeString(test.nonce)
	plaintext := grand.GetRandom(0)
	ad := []byte(test.ad)
	ct1 := gcm.Seal(nil, nonce, plaintext, ad)

	ct2 := make([]byte, 0, 116)
	err = gcm.EncryptInit(nonce)
	if err != nil {
		t.Fatal(err)
	}
	gcm.SpecifyADD(ad)
	for i := range plaintext {
		ct2, err = gcm.EncryptUpdate(ct2, []byte{plaintext[i]})
		if err != nil {
			t.Fatal(err)
		}
	}
	ct2, err = gcm.EncryptFinal(ct2)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(ct1, ct2) {
		t.Errorf("Seal failed")
	}

	if err = gcm.DecryptInit(nonce); err != nil {
		t.Fatal(err)
	}
	gcm.SpecifyADD(ad)

	decrypted := make([]byte, 0, len(plaintext))
	for i := range ct1 {
		decrypted, err = gcm.DecryptUpdate(decrypted, []byte{ct1[i]})
		if err != nil {
			t.Fatal(err)
		}
	}
	decrypted, err = gcm.DecryptFinal(decrypted)
	if !bytes.Equal(decrypted, plaintext) {
		t.Fatal("plaintext unequal decrypted plaintext")
	}
	if err != nil {
		t.Fatal("auth failed")
	}
}

func FuzzSm4Gcm(f *testing.F) {
	nonce := grand.GetRandom(12)
	key := grand.GetRandom(16)

	block, _ := sm4.NewCipher(key)
	gcm, err := blockmode.NewGCM(blockmode.Wrap(block))
	if err != nil {
		f.Fatal(err)
	}
	stdgcm, _ := sm4.NewGCM(key)

	f.Add([]byte{}, []byte{})
	f.Fuzz(func(t *testing.T, plaintext, ad []byte) {
		stdct := stdgcm.Seal(nil, nonce, plaintext, ad)
		ct1 := gcm.Seal(nil, nonce, plaintext, ad)
		if !bytes.Equal(ct1, stdct) {
			t.Errorf("Seal failed")
		}

		if err := gcm.EncryptInit(nonce); err != nil {
			t.Fatal(err)
		}
		// if additional data is empty, then call
		// gcm.SpecifyADD(nil)
		gcm.SpecifyADD(ad)

		ct2, err := gcm.EncryptUpdate(nil, plaintext[:len(plaintext)/2])
		if err != nil {
			t.Fatal(err)
		}
		for _, p := range plaintext[len(plaintext)/2:] {
			if ct2, err = gcm.EncryptUpdate(ct2, []byte{p}); err != nil {
				t.Fatal(err)
			}
		}
		ct2, err = gcm.EncryptFinal(ct2)
		if err != nil {
			t.Fatal(err)
		}

		if !bytes.Equal(ct2, stdct) {
			t.Errorf("Encrypt failed")
		}

		pt1, err1 := gcm.Open(nil, nonce, ct1, ad)
		if err1 != nil {
			t.Errorf("Open faile: %v\n", err1)
		}
		if err := gcm.DecryptInit(nonce); err != nil {
			t.Fatal(err)
		}
		gcm.SpecifyADD(ad)

		pt2, err := gcm.DecryptUpdate(nil, ct2[:len(ct2)/2])
		if err != nil {
			t.Fatal(err)
		}
		for _, p := range ct2[len(ct2)/2:] {
			if pt2, err = gcm.DecryptUpdate(pt2, []byte{p}); err != nil {
				t.Fatal(err)
			}
		}
		pt2, err = gcm.DecryptFinal(pt2)
		if err != nil {
			t.Fatal(err)
		}
		if !bytes.Equal(pt2, plaintext) {
			t.Errorf("Decrypt failed")
		}

		if !bytes.Equal(pt1, plaintext) {
			t.Errorf("Open failed")
		}
	})
}