package sm9

// fixed data test

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"xdx.jelly/xgcl/grand"
)

func TestKey(t *testing.T) {
	uid := []byte("Alice")
	b, _ := hex.DecodeString("000130E78459D78545CB54C587E02CF480CE0B66340F319F348A1D5B1F2DC5F4")
	ks, pubs, _ := GenerateMastSignPrivateKey(bytes.NewReader(b))

	if hex.EncodeToString(pubs.Bytes()) != "9f64080b3084f733e48aff4b41b565011ce0711c5e392cfb0ab1b6791b94c40829dba116152d1f786ce843ed24a3b573414d2177386a92dd8f14d65696ea5e3269850938abea0112b57329f447e3a0cbad3e2fdb1a77f335e89e1408d0ef1c2541e00a53dda532da1a7ce027b7a46f741006e85f5cdff0730e75c05fb4e3216d" {
		t.Fatal()
		return
	}

	ds, err := ks.GenerateUserSignKey(uid)
	if err != nil {
		t.Log(err)
		t.Fatal()
		return
	}

	if hex.EncodeToString(ds.Bytes()) != "a5702f05cf1315305e2d6eb64b0deb923db1a0bcf0caff90523ac8754aa6982078559a844411f9825c109f5ee3f52d720dd01785392a727bb1556952b2b013d3" {
		t.Fatal()
		return
	}

	uid = []byte("Bob")
	b, _ = hex.DecodeString("0001EDEE3778F441F8DEA3D9FA0ACC4E07EE36C93F9A08618AF4AD85CEDE1C22")
	ke, pube, _ := GenerateMastEncPrivateKey(bytes.NewReader(b))

	if hex.EncodeToString(pube.Bytes()) != "787ed7b8a51f3ab84e0a66003f32da5c720b17eca7137d39abc66e3c80a892ff769de61791e5adc4b9ff85a31354900b202871279a8c49dc3f220f644c57a7b1" {
		t.Fatal()
		return
	}
	de, _ := ke.GenerateUserEncKey(uid)
	buf := de.Bytes()
	if err := de.SetBytes(buf); err != nil {
		t.Fatal(err)
	}

	if hex.EncodeToString(de.Bytes()) != "94736acd2c8c8796cc4785e938301a139a059d3537b6414140b2d31eecf41683115bae85f5d8bc6c3dbd9e5342979acccf3c2f4f28420b1cb4f8c0b59a19b1587aa5e47570da7600cd760a0cf7beaf71c447f3844753fe74fa7ba92ca7d3b55f27538a62e7f7bfb51dce08704796d94c9d56734f119ea44732b50e31cdeb75c1" {
		t.Fatal()
		return
	}
}

func TestSign(t *testing.T) {
	uid := []byte("Alice")
	msg := []byte("Chinese IBS standard")
	rnd, _ := hex.DecodeString("000130E78459D78545CB54C587E02CF480CE0B66340F319F348A1D5B1F2DC5F4")
	ks, pubs, _ := GenerateMastSignPrivateKey(bytes.NewReader(rnd))

	ds, _ := ks.GenerateUserSignKey(uid)
	rnd, _ = hex.DecodeString("00033C8616B06704813203DFD00965022ED15975C662337AED648835DC4B1CBE")
	_, _ = Sign(msg, ds, pubs, rnd)
	// Test Issue 5, Sign changed ds.
	signature, _ := Sign(msg, ds, pubs, rnd)

	s, err := signature.MarshalBinary()
	if err != nil {
		t.Log("Signature.MarshalBinary failed: ", err)
	}
	if hex.EncodeToString(s) !=
		"823c4b21e4bd2dfe1ed92c606653e996668563152fc33f55d7bfbb9bd9705adb"+
			"73bf96923ce58b6ad0e13e9643a406d8eb98417c50ef1b29cef9adb48b6d598c"+
			"856712f1c2e0968ab7769f42a99586aed139d5b8b3e15891827cc2aced9baa05" {
		t.Fatal()
	}
	if !Verify(signature, uid, msg, pubs) {
		t.Fatal()
	}
}

func TestKeyExchange(t *testing.T) {
	// FIXME test the fixed data
	idA := []byte("Sponsor")
	idB := []byte("Responsor")

	keA, pubeA, _ := GenerateMastEncPrivateKey(rand.Reader)
	deA, _ := keA.GenerateUserEncKey(idA)
	// A and B must under a same KGC's master key
	keB := keA
	// keB := GenMastEncPrivateKey(nil)
	pubeB := pubeA
	// pubeB := GenMastEncPublicKey(keB)
	deB, _ := keB.GenerateUserEncKey(idB)

	// 使用默认id传入nil或GetDefaultID
	alice := NewSponsor(idA, deA)
	bob := NewResponsor(idB, deB)

	// NEVER FORGET CLEAR
	// defer keA.Clear()
	// defer keB.Clear()
	// defer deA.Clear()
	// defer deB.Clear()
	defer alice.Clear()
	defer bob.Clear()

	// 多次密钥交换是可以重复使用Sponsor和Responsor的-只要其私钥和id不变
	for keylen := 1; keylen < 256; keylen++ {
		// t.Log("key exchange test for key length = ", keylen)

		// key exchange
		tempKeyOfSponsor, _ := alice.GenerateAgreementData(idB, pubeB, nil)
		// t.Log("11", tempKeyOfSponsor, err)
		keyOfResponsor, tempKeyOfResponsor, _ := bob.GenerateAgreementDataAndKey(idA, pubeA, tempKeyOfSponsor, keylen, nil)
		// t.Log("tempKeyOfResponsor", tempKeyOfResponsor)
		// return

		keyOfSponsor, _ := alice.GenerateKey(tempKeyOfResponsor, keylen)

		// check if OK
		// printlen := 32
		// if keylen < printlen {
		// 	printlen = keylen
		// }
		//t.Log("responsor: ", hex.EncodeToString(keyOfResponsor[:printlen]), "...")
		//t.Log("sponsor  : ", hex.EncodeToString(keyOfSponsor[:printlen]), "...")
		if !bytes.Equal(keyOfResponsor, keyOfSponsor) {
			t.Log()
			t.Fail()
			return
		}
		// t.Log("OK\n")

	}
}

func TestKeyExchange2(t *testing.T) {
	idA := []byte("Sponsor")
	idB := []byte("Responsor")

	ke, pube, _ := GenerateMastEncPrivateKey(rand.Reader)
	deA, _ := ke.GenerateUserEncKey(idA)
	deB, _ := ke.GenerateUserEncKey(idB)

	for keylen := 0; keylen < 256; keylen++ {
		ra, Ra, err := GenerateAgreementData(idB, pube, grand.Reader)
		if err != nil {
			t.Fatal()
		}
		key2, Rb, err := GenerateAgreementDataAndKey(idA, idB, Ra, deB, pube, keylen, grand.Reader)
		if err != nil {
			t.Fatal()
		}
		key1, err := GenerateKey(idA, idB, Ra, Rb, ra, deA, pube, keylen)
		if err != nil {
			t.Fatal()
		}

		if !bytes.Equal(key1, key2) {
			t.Fatal()
		}

	}
}

func TestKeyEncap(t *testing.T) {
	r, _ := hex.DecodeString("0001EDEE3778F441F8DEA3D9FA0ACC4E07EE36C93F9A08618AF4AD85CEDE1C22")
	ke, pube, _ := GenerateMastEncPrivateKey(bytes.NewReader(r))
	id := []byte("Bob")

	r, _ = hex.DecodeString("000074015F8489C01EF4270456F9E6475BFB602BDE7F33FD482AB4E3684A6722")

	klen := 32
	keypackage, key, err := KeyEncapsulate(id, klen, pube, r)
	if err != nil {
		t.Fatal(err)
	}
	if hex.EncodeToString(key) != "4ff5cf86d2ad40c8f4bac98d76abdbde0c0e2f0a829d3f911ef5b2bce0695480" {
		t.Fatal()
	}

	target, _ := hex.DecodeString("1edee2c3f465914491de44cefb2cb434ab02c308d9dc5e2067b4fed5aaac8a0f1c9b4c435eca35ab83bb734174c0f78fde81a53374aff3b3602bbc5e37be9a4c")
	if !bytes.Equal(keypackage.G1.Marshal(), target) {
		t.Fatal()
	}

	de, _ := ke.GenerateUserEncKey(id)
	uncapKey, _ := KeyDecapsulate(id, keypackage, klen, de)
	if hex.EncodeToString(uncapKey) != "4ff5cf86d2ad40c8f4bac98d76abdbde0c0e2f0a829d3f911ef5b2bce0695480" {
		t.Log(hex.EncodeToString(uncapKey))
		t.Log(hex.EncodeToString(key))
		t.Fatal()
	}
}

func TestEncryptionECB(t *testing.T) {
	id := []byte("Bob")
	rnd, _ := hex.DecodeString("0001EDEE3778F441F8DEA3D9FA0ACC4E07EE36C93F9A08618AF4AD85CEDE1C22")
	ke, pube, _ := GenerateMastEncPrivateKey(bytes.NewReader(rnd))

	msg := []byte("Chinese IBE standard")
	msg = append(msg, []byte{0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c, 0x0c}...)
	de, _ := ke.GenerateUserEncKey(id)
	rnd, _ = hex.DecodeString("0000AAC0541779C8FC45E3E2CB25C12B5D2576B2129AE8BB5EE2CBE5EC9E785C")
	c, err := EncryptionSm4ECB(id, msg, pube, rnd)
	if err != nil {
		t.Fatal()
	}
	encType := make([]byte, 4)
	binary.BigEndian.PutUint32(encType, uint32(EncTypeSM4ECB))
	if data, err := c.MarshalBinary(); err != nil || hex.EncodeToString(data) != hex.EncodeToString(encType)+"2445471164490618e1ee20528ff1d545b0f14c8bcaa44544f03dab5dac07d8ff42ffca97d57cddc05ea405f2e586feb3a6930715532b8000759f13059ed59ac0fd3c98dd92c44c68332675a370cceede31e0c5cd209c257601149d12b394a2be00000020e05b6fac6f11b965268c994f00dba7a8bb00fd60583546cbdf4649250863f10a" {
		t.Log("Got :", hex.EncodeToString(data))
		t.Log("Want:", "000000022445471164490618e1ee20528ff1d545b0f14c8bcaa44544f03dab5dac07d8ff42ffca97d57cddc05ea405f2e586feb3a6930715532b8000759f13059ed59ac0fd3c98dd92c44c68332675a370cceede31e0c5cd209c257601149d12b394a2be00000020e05b6fac6f11b965268c994f00dba7a8bb00fd60583546cbdf4649250863f10a")
		t.Fatal()
	}

	plain, err := DecryptionSm4ECB(id, c, de)
	if !bytes.Equal(plain, msg) || err != nil {
		t.Fatal()
	}

	data, _ := c.MarshalBinary()
	c1 := NewCipher()
	if err = c1.UnmarshalBinary(data); err != nil {
		t.Fatal()
	}
	data2, _ := c1.MarshalBinary()
	if !bytes.Equal(data, data2) {
		t.Fatal()
	}

}

func TestEncryptionKDF(t *testing.T) {
	id := []byte("Bob")
	rnd, _ := hex.DecodeString("0001EDEE3778F441F8DEA3D9FA0ACC4E07EE36C93F9A08618AF4AD85CEDE1C22")
	ke, pube, _ := GenerateMastEncPrivateKey(bytes.NewReader(rnd))

	msg := []byte("Chinese IBE standard")
	de, _ := GenerateUserEncKey(id, ke)
	rnd, _ = hex.DecodeString("0000AAC0541779C8FC45E3E2CB25C12B5D2576B2129AE8BB5EE2CBE5EC9E785C")

	c, err := Encrypt(EncTypeKDF, pube, id, msg, bytes.NewReader(rnd), nil)
	if err != nil {
		t.Fatal(err)
	}
	encType := make([]byte, 4)
	binary.BigEndian.PutUint32(encType, uint32(EncTypeKDF))
	if data, err := c.MarshalBinary(); err != nil || hex.EncodeToString(data) != hex.EncodeToString(encType)+"2445471164490618e1ee20528ff1d545b0f14c8bcaa44544f03dab5dac07d8ff42ffca97d57cddc05ea405f2e586feb3a6930715532b8000759f13059ed59ac0ba672387bcd6de5016a158a52bb2e7fc429197bcab70b25afee37a2b9db9f367000000141b5f5b0e951489682f3e64e1378cdd5da9513b1c" {
		t.Fatal()
	}

	plain, err := Decrypt(id, c, de)
	if !bytes.Equal(plain, msg) || err != nil {
		t.Fatal()
	}

	data, _ := c.MarshalBinary()
	c1 := &Cipher{}
	if err = c1.UnmarshalBinary(data); err != nil {
		t.Fatal()
	}
	data2, _ := c1.MarshalBinary()
	if !bytes.Equal(data, data2) {
		t.Fatal()
	}

}

func TestGenMastKey(t *testing.T) {
	for i := 1; i < 1000; i++ {
		_, _, err := GenerateMastEncPrivateKey(bytes.NewReader(grand.GetRandom(i)))
		if err != nil {
			t.Fatal()
		}

		_, _, err = GenerateMastSignPrivateKey(bytes.NewReader(grand.GetRandom(i)))
		if err != nil {
			t.Fatal()
		}
	}
}

func TestSign1(t *testing.T) {
	uid, _ := hex.DecodeString("5573657231")
	msg, _ := hex.DecodeString("5152536162636471")

	d, _ := hex.DecodeString("530569D472BD8A263AF99F6A3DE7FAA807E2B1094A6DCB98EAFE7E045B64F7EF" +
		"88201ADF41654F340A796F113D0E885BA70D088B076F14831B4553EB0105B2E3" +
		"0FFFDC648D8BD259F27BB90C9C09E6E8FAEDD46D1B017FF3CEFA6FC914EFE8D5" +
		"3137D9A8731FFDA255E519CE3340DC1D0ED5F0273258E441B994EC955B497961")
	pubs, _ := new(MastSignPublicKey).SetBytes(d)

	d, _ = hex.DecodeString("1740268704F86C31B641287A1B296087DFDB43C5FB00AB687D03059E54334583" +
		"08B1B3DBD5447690DC8EC3967D804927AF5C6BD36ED058EB9728BDAF69B117C7" +
		"359DA014F0C619F9729AE1F2B62C567B1DCF5E1885DA322B019DF1F04F3CDFF8")
	signature := new(Signature)
	signature.SetBytes(d)

	if !Verify(signature, uid, msg, pubs) {
		t.Fatal()
	}
}

func TestEncryption1(t *testing.T) {
	id, err := hex.DecodeString("5573657231")
	assert.Nil(t, err)
	de := new(UserEncKey)
	b, err := hex.DecodeString("5DB6FAEF0E0C010C20F35341C68DA9111BFB8582C60F9B07F08E3D9F7929D14F847DA2D9F02ACA4B324827A2B54601CE2C26D9693F6A3018984925557F11653451D82AE921FD2C9D64D8D3A824DF1B6753C71234A5420E1BAC456919B88E84EB8D7B9FA5788AEF0B9BAD1BECC5FAA9619BA58F2C97343235193A24857EF9F0A3")
	assert.Nil(t, err)
	de.Unmarshal(b)

	b, err = hex.DecodeString("23E2C8C3C2267865198A6C1062D7703A0745D9AB86CCB3E0CB0E9B8D9D9D25C16AE46F6C174945020BDD8F679024B6179462BE4A9AE648999D909432AB7EEEFCCD34790782D43D80F0713C51F74C9B41E55F6F10496AD18D6EBC85626F5236961643E90D0C9DFB9AC80B7BA9B69F35D7")
	assert.Nil(t, err)
	c := &Cipher{EncType: 1}
	c.C1.Unmarshal(b[:64])
	copy(c.H[:], b[64:96])
	c.C = append([]byte{}, b[96:]...)

	plain, err := Decrypt(id, c, de)
	assert.Nil(t, err)

	fmt.Printf("%x", plain)
}

func TestEncryption2(t *testing.T) {
	id, err := hex.DecodeString("426f62")
	assert.Nil(t, err)
	de := new(UserEncKey)
	b, err := hex.DecodeString("94736acd2c8c8796cc4785e938301a139a059d3537b6414140b2d31eecf41683115bae85f5d8bc6c3dbd9e5342979acccf3c2f4f28420b1cb4f8c0b59a19b1587aa5e47570da7600cd760a0cf7beaf71c447f3844753fe74fa7ba92ca7d3b55f27538a62e7f7bfb51dce08704796d94c9d56734f119ea44732b50e31cdeb75c1")
	assert.Nil(t, err)
	de.Unmarshal(b)

	b, err = hex.DecodeString("2445471164490618e1ee20528ff1d545b0f14c8bcaa44544f03dab5dac07d8ff42ffca97d57cddc05ea405f2e586feb3a6930715532b8000759f13059ed59ac0ba672387bcd6de5016a158a52bb2e7fc429197bcab70b25afee37a2b9db9f3671b5f5b0e951489682f3e64e1378cdd5da9513b1c")
	assert.Nil(t, err)
	c := &Cipher{EncType: 0}
	c.C1.Unmarshal(b[:64])
	copy(c.H[:], b[64:96])
	c.C = append([]byte{}, b[96:]...)

	plain, err := Decrypt(id, c, de)
	assert.Nil(t, err)

	fmt.Printf("%x", plain)
}