package sm9m

import (
	"fmt"
	"testing"

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

func TestAuditSign(t *testing.T) {
	uid := []byte("Alice")

	ks, pubs, err := sm9.GenerateMastSignPrivateKey(drng.SM3Rng) // use SM3 DRNG
	assert.Nil(t, err)
	fmt.Printf("签名主公私钥: %x\n", ks.Bytes())
	fmt.Printf("签名主公钥: %x\n", pubs.Bytes())

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

	clientSignKey, serverSignKey, err := GenerateSignKeys(ds, drng.SM3Rng)
	assert.Nil(t, err)
	fmt.Printf("签名公钥: %x\n", uid)
	fmt.Printf("客户端密钥: %x\n", clientSignKey.p.Marshal())
	fmt.Printf("服务端密钥: %x\n", serverSignKey.k.Bytes())

	// SM9签名输入SM3摘要
	msg := grand.GetRandom(32)
	fmt.Printf("签名原文: %x\n", msg)
	hash := sm3.Sum(msg)

	// 客户端第一步
	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)
	fmt.Printf("签名结果: %x\n", sig.Bytes())

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

func TestAuditDec(t *testing.T) {

	uid := []byte("Alice")

	ke, pube, err := sm9.GenerateMastEncPrivateKey(drng.SM3Rng) // use SM3 DRNG
	assert.Nil(t, err)
	fmt.Printf("加密主公私钥: %x\n", ke.Bytes())
	fmt.Printf("加密主公钥: %x\n", pube.Bytes())

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

	clientEncKey, serverEncKey, err := GenerateEncKeys(de, drng.SM3Rng)
	assert.Nil(t, err)
	fmt.Printf("签名公钥: %x\n", uid)
	fmt.Printf("客户端密钥: %x\n", clientEncKey.p.Marshal())
	fmt.Printf("服务端密钥: %x\n", serverEncKey.p.Marshal())
	{

		fmt.Println("===== SM9协同解密 =====")
		msg := grand.GetRandom(32)
		fmt.Printf("原文: %x\n", msg)
		cipher, err := sm9.Encrypt(sm9.EncTypeKDF, pube, uid, msg, drng.SM3Rng, nil)
		assert.Nil(t, err)
		b, _ := cipher.MarshalBinary()
		fmt.Printf("加密密文: %x\n", b)

		// 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(uid, clientEncKey, serverParam)
		assert.Nil(t, err)

		assert.Equal(t, msg, msg1)
		fmt.Printf("解密明文: %x\n", msg1)
	}

	fmt.Println("===== SM9协同解封装 =====")

	{
		keylen := 48
		pack, key0, err := sm9.KeyEncapsulate(uid, keylen, pube, drng.SM3Rng)
		assert.Nil(t, err)
		fmt.Printf("封装密文: %x\n", pack.Marshal())
		fmt.Printf("封装密钥: %x\n", key0)

		// 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(uid, keylen, clientEncKey, serverParam)
		assert.Nil(t, err)

		fmt.Printf("解封装密钥: %x\n", key1)

		assert.Equal(t, key0, key1)
	}
}