package sm9m

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"xdx.jelly/xgcl/grand/drng"
	"xdx.jelly/xgcl/sm/sm3"
	"xdx.jelly/xgcl/sm/sm9"
)

func TestSignKeyGen(t *testing.T) {
	uid := []byte("Alice")
	ks, pubs, err := sm9.GenerateMastSignPrivateKey(drng.SM3Rng) // use SM3 DRNG
	assert.Nil(t, err)
	_ = pubs

	ds, err := ks.GenerateUserSignKey(uid)
	assert.Nil(t, err)

	clientSignKey, serverSignKey, err := GenerateSignKeys(ds, drng.SM3Rng)
	assert.Nil(t, err)

	// Check generated keys are valid
	assert.True(t, clientSignKey.p.ScalarMult(&clientSignKey.p, &serverSignKey.k).Equal(&ds.G1))
}

func TestEncKeyGen(t *testing.T) {
	uid := []byte("Alice")
	ke, _, err := sm9.GenerateMastEncPrivateKey(drng.SM3Rng) // use SM3 DRNG
	assert.Nil(t, err)

	de, err := ke.GenerateUserEncKey(uid)
	assert.Nil(t, err)

	clientEncKey, serverEncKey, err := GenerateEncKeys(de, drng.SM3Rng)
	assert.Nil(t, err)

	// Check generated keys are valid
	assert.True(t, serverEncKey.p.Add(&serverEncKey.p, &clientEncKey.p).Equal(&de.G2))
}

func TestSign(t *testing.T) {
	uid := []byte("Alice")
	ks, pubs, err := sm9.GenerateMastSignPrivateKey(drng.SM3Rng) // use SM3 DRNG
	assert.Nil(t, err)

	ds, err := ks.GenerateUserSignKey(uid)
	assert.Nil(t, err)

	clientSignKey, serverSignKey, err := GenerateSignKeys(ds, drng.SM3Rng)
	assert.Nil(t, err)

	// SM9签名输入SM3摘要
	hash := sm3.Sum([]byte("abc"))

	// 客户端第一步
	clientCTX := &ClientSignContext{}
	param, err := clientCTX.GenerateParam(clientSignKey, drng.SM3Rng)
	assert.Nil(t, err)
	b, err := param.MarshalBinary()
	assert.Nil(t, err)

	// 服务端
	param = &ClientSignParam{}
	assert.Nil(t, param.UnmarshalBinary(b))
	serverParam, err := ServerSign(hash[:], param, serverSignKey, pubs, drng.SM3Rng)
	assert.Nil(t, err)
	b, err = serverParam.MarshalBinary()
	assert.Nil(t, err)

	// 客户端第二步
	serverParam = &ServerSignParam{}
	assert.Nil(t, serverParam.UnmarshalBinary(b))
	sig, err := clientCTX.GenerateSignature(serverParam)
	assert.Nil(t, err)

	assert.True(t, sm9.Verify(sig, uid, hash[:], pubs))
}

func TestDecryption(t *testing.T) {
	id := []byte("Alice")
	ke, pube, err := sm9.GenerateMastEncPrivateKey(drng.SM3Rng) // use SM3 DRNG
	assert.Nil(t, err)

	de, err := ke.GenerateUserEncKey(id)
	assert.Nil(t, err)

	msg := []byte("Chinese IBE standard")
	cipher, err := sm9.Encrypt(sm9.EncTypeKDF, pube, id, msg, drng.SM3Rng, nil)
	assert.Nil(t, err)

	clientEncKey, serverEncKey, err := GenerateEncKeys(de, drng.SM3Rng)
	assert.Nil(t, err)

	// 1
	clientCTX := &ClientDecContext{}
	clientParam, err := clientCTX.GenerateParam(cipher)
	assert.Nil(t, err)
	b, err := clientParam.MarshalBinary()
	assert.Nil(t, err)

	// 2
	clientParam = &ClientDecParam{}
	assert.Nil(t, clientParam.UnmarshalBinary(b))
	serverParam, err := ServerDecrypt(clientParam, serverEncKey)
	assert.Nil(t, err)
	b, err = serverParam.MarshalBinary()
	assert.Nil(t, err)

	// 3
	serverParam = &ServerDecParam{}
	assert.Nil(t, serverParam.UnmarshalBinary(b))
	msg1, err := clientCTX.DecryptFinal(id, clientEncKey, serverParam)
	assert.Nil(t, err)

	assert.Equal(t, msg, msg1)
}

func TestDecapsulate(t *testing.T) {
	id := []byte("Alice")
	ke, pube, err := sm9.GenerateMastEncPrivateKey(drng.SM3Rng) // use SM3 DRNG
	assert.Nil(t, err)

	de, err := ke.GenerateUserEncKey(id)
	assert.Nil(t, err)

	clientEncKey, serverEncKey, err := GenerateEncKeys(de, drng.SM3Rng)
	assert.Nil(t, err)

	keylen := 48
	pack, key0, err := sm9.KeyEncapsulate(id, keylen, pube, drng.SM3Rng)
	assert.Nil(t, err)

	// 1
	clientCTX := &ClientDecapsulateContext{}
	clientParam, err := clientCTX.GenerateParam(pack)
	assert.Nil(t, err)
	b, err := clientParam.MarshalBinary()
	assert.Nil(t, err)

	// 2
	clientParam = &ClientDecapsulateParam{}
	assert.Nil(t, clientParam.UnmarshalBinary(b))
	serverParam, err := ServerDecapsulate(clientParam, serverEncKey)
	assert.Nil(t, err)
	b, err = serverParam.MarshalBinary()
	assert.Nil(t, err)

	// 3
	serverParam = &ServerDecapsulateParam{}
	assert.Nil(t, serverParam.UnmarshalBinary(b))
	key1, err := clientCTX.DecapsulateFinal(id, keylen, clientEncKey, serverParam)
	assert.Nil(t, err)

	assert.Equal(t, key0, key1)
}