package sdf_test

import (
	"bytes"
	"testing"

	"xdx.jelly/xgcl/api/common"
	"xdx.jelly/xgcl/api/sdf"
	"xdx.jelly/xgcl/identifier"
	"xdx.jelly/xgcl/sm/sm2"
)

var gsdf sdf.Sdfable
var index uint32 = 1
var peerIndex uint32 = 2

func init() {
	SDF := &sdf.SdfNoLock{}
	for idx := uint32(1); idx < 10000; idx++ {
		refKey := &common.ECCrefPrivateKey{Bits: 256}
		SDF.SDF_GenerateRandom(refKey.K[32:])

		var err error
		err = SDF.ImportSm2KeyAtIndex(idx, sdf.KeyTypeSm2Sign, refKey)
		if err != nil {
			panic(err)
		}
		err = SDF.ImportSm2KeyAtIndex(idx, sdf.KeyTypeSm2Enc, refKey)
		if err != nil {
			panic(err)
		}
	}
	for idx := uint32(0); idx < 10000; idx++ {
		SDF.GenerateKekAtIndex(idx)
	}

	gsdf = SDF
}

func TestDevInfo(t *testing.T) {
	_, err := gsdf.SDF_GetDeviceInfo()
	if err != nil {
		t.Fatal(err)
	}
}

func TestGenerateRandom(t *testing.T) {
	buf := make([]byte, 32)
	if n, err := gsdf.SDF_GenerateRandom(buf); err != nil || n < uint32(len(buf)) {
		if err != nil {
			t.Fatal(err)
		} else {
			t.Fatalf("Generate %d random bytes, desired %d\n", n, len(buf))
		}
	}
}

func TestSign(t *testing.T) {
	buf := make([]byte, 32)

	sig, err := gsdf.SDF_InternalSign_ECC(index, buf)
	if err != nil {
		t.Fatal(err)
	}

	err = gsdf.SDF_InternalVerify_ECC(index, buf, sig)
	if err != nil {
		t.Fatal(err)
	}
}

func TestGenerateKey_ECC(t *testing.T) {
	gsdf.SDF_OpenSession()
	defer gsdf.SDF_CloseSession()
	encKey, handle1, err := gsdf.SDF_GenerateKeyWithIPK_ECC(index, 128)
	if err != nil {
		t.Fatal(err)
	}
	defer gsdf.SDF_DestroyKey(handle1)

	handle2, err := gsdf.SDF_ImportKeyWithISK_ECC(index, encKey)
	if err != nil {
		t.Fatal(err)
	}
	defer gsdf.SDF_DestroyKey(handle2)

	var data [128]byte
	var iv [16]byte
	encData := []byte{}
	decData := []byte{}
	_, _ = gsdf.SDF_GenerateRandom(data[:])
	_, _ = gsdf.SDF_GenerateRandom(iv[:])

	err = gsdf.SDF_Encrypt(handle1, identifier.SGDSM4CBC, append([]byte{}, iv[:]...), data[:], &encData)
	if err != nil {
		t.Fatal(err)
	}

	err = gsdf.SDF_Decrypt(handle2, identifier.SGDSM4CBC, append([]byte{}, iv[:]...), encData, &decData)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Compare(data[:], decData) != 0 {
		t.Fatal("Compare unequal")
	}
}

func TestGenerateKey_KEK(t *testing.T) {
	gsdf.SDF_OpenSession()
	defer gsdf.SDF_CloseSession()
	encKey, handle1, err := gsdf.SDF_GenerateKeyWithKEK(128, identifier.SGDSM4ECB, 0)
	if err != nil {
		t.Fatal(err)
	}
	defer gsdf.SDF_DestroyKey(handle1)

	handle2, err := gsdf.SDF_ImportKeyWithKEK(identifier.SGDSM4ECB, 0, encKey)
	if err != nil {
		t.Fatal(err)
	}
	defer gsdf.SDF_DestroyKey(handle2)

	var data [128]byte
	var iv [16]byte
	encData := []byte{}
	decData := []byte{}
	_, _ = gsdf.SDF_GenerateRandom(data[:])
	_, _ = gsdf.SDF_GenerateRandom(iv[:])

	err = gsdf.SDF_Encrypt(handle1, identifier.SGDSM4CBC, append([]byte{}, iv[:]...), data[:], &encData)
	if err != nil {
		t.Fatal(err)
	}

	err = gsdf.SDF_Decrypt(handle2, identifier.SGDSM4CBC, append([]byte{}, iv[:]...), encData, &decData)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Compare(data[:], decData) != 0 {
		t.Fatal("Compare unequal")
	}
}

func TestExportPublicKey_ECC(t *testing.T) {
	gsdf.SDF_OpenSession()
	defer gsdf.SDF_CloseSession()

	pk, err := gsdf.SDF_ExportSignPublicKey_ECC(index)
	if err != nil {
		t.Fatal(err)
	}
	e := make([]byte, 32)
	sig, err := gsdf.SDF_InternalSign_ECC(index, e)
	if err != nil {
		t.Fatal(err)
	}

	err = gsdf.SDF_ExternalVerify_ECC(identifier.SGDSM21, pk, e, sig)
	if err != nil {
		t.Fatal(err)
	}
}

func TestGenerateKey_SSL(t *testing.T) {

	preMasterSecret := make([]byte, 48)
	preMasterSecret[0] = 1
	preMasterSecret[1] = 1
	gsdf.SDF_GenerateRandom(preMasterSecret[2:])

	var clientServerRandom [32 * 2]byte
	gsdf.SDF_GenerateRandom(clientServerRandom[:])

	pucPublicKey, err := gsdf.SDF_ExportEncPublicKey_ECC(index)

	pucKeyClientMac, phKeyHandleClientMac, pucKeyServerMac, phKeyHandleServerMac, pucKeyClientEnc, phKeyHandleClientEnc, pucKeyServerEnc, phKeyHandleServerEnc,
		clientWriteIV, serverWriteIV, err := gsdf.SDF_GenerateKeywithEPK_SSL(preMasterSecret, clientServerRandom[:32], clientServerRandom[32:], identifier.SGDSM3, identifier.SGDSM23, pucPublicKey, 128, 128, 128, 128)
	if err != nil {
		t.Fatal(err)
	}
	keyHandleClientMac, err := gsdf.SDF_ImportKeyWithISK_ECC(index, pucKeyClientMac)
	keyHandleServerMac, err := gsdf.SDF_ImportKeyWithISK_ECC(index, pucKeyServerMac)
	keyHandleClientEnc, err := gsdf.SDF_ImportKeyWithISK_ECC(index, pucKeyClientEnc)
	keyHandleServerEnc, err := gsdf.SDF_ImportKeyWithISK_ECC(index, pucKeyServerEnc)

	msg := make([]byte, 128)
	gsdf.SDF_GenerateRandom(msg)

	clientMac1, err := gsdf.SDF_CalculateMAC(keyHandleClientMac, identifier.SGDSM3, nil, msg)
	clientMac2, err := gsdf.SDF_CalculateMAC(phKeyHandleClientMac, identifier.SGDSM3, nil, msg)
	if bytes.Compare(clientMac1, clientMac2) != 0 {
		t.Fatal("client mac error")
	}

	serverMac1, err := gsdf.SDF_CalculateMAC(keyHandleServerMac, identifier.SGDSM3, nil, msg)
	serverMac2, err := gsdf.SDF_CalculateMAC(phKeyHandleServerMac, identifier.SGDSM3, nil, msg)
	if bytes.Compare(serverMac1, serverMac2) != 0 {
		t.Fatal("server mac error")
	}

	encMsg := make([]byte, len(msg))
	decMsg := make([]byte, len(msg))
	err = gsdf.SDF_Encrypt(keyHandleClientEnc, identifier.SGDSM4CBC, append([]byte{}, clientWriteIV...), msg, &encMsg)
	err = gsdf.SDF_Decrypt(phKeyHandleClientEnc, identifier.SGDSM4CBC, append([]byte{}, clientWriteIV...), encMsg, &decMsg)
	if bytes.Compare(msg, decMsg) != 0 {
		t.Fatal("client enc error")
	}

	err = gsdf.SDF_Encrypt(keyHandleServerEnc, identifier.SGDSM4CBC, append([]byte{}, serverWriteIV...), msg, &encMsg)
	err = gsdf.SDF_Decrypt(phKeyHandleServerEnc, identifier.SGDSM4CBC, append([]byte{}, serverWriteIV...), encMsg, &decMsg)
	if bytes.Compare(msg, decMsg) != 0 {
		t.Fatal("server enc error")
	}

}

func TestGenerateKey_ECDHE_SSL(t *testing.T) {
	//TODO 预主密钥48字节，前两字节应该是版本号。生成46字节的协商密钥
	pucSponsorPublicKey, pucSponsorTmpPublicKey, phAgreementHandle, err := gsdf.SDF_GenerateAgreementDataWithECC(index, 46*8, sm2.GetDefaultID())
	if err != nil {
		t.Fatal(err)
	}

	pucResponsePublicKey, pucResponseTmpPublicKey, phKeyHandle2, err := gsdf.SDF_GenerateAgreementDataAndKeyWithECC(peerIndex, 46*8, sm2.GetDefaultID(), sm2.GetDefaultID(), pucSponsorPublicKey, pucSponsorTmpPublicKey)
	if err != nil {
		t.Fatal(err)
	}

	phKeyHandle1, err := gsdf.SDF_GenerateKeyWithECC(sm2.GetDefaultID(), pucResponsePublicKey, pucResponseTmpPublicKey, phAgreementHandle)
	if err != nil {
		t.Fatal(err)
	}

	var clientServerRandom [32 * 2]byte
	gsdf.SDF_GenerateRandom(clientServerRandom[:])
	pucClientRandom := clientServerRandom[:32]
	pucServerRandom := clientServerRandom[:32]

	phKeyHandleClientMac1, phKeyHandleServerMac1, phKeyHandleClientEnc1, phKeyHandleServerEnc1, clientWriteIV1, serverWriteIV1, err := gsdf.SDF_GenerateKeywithECDHE_SSL(phKeyHandle1, pucClientRandom, pucServerRandom, identifier.SGDSM3, 128, 128, 128, 128)
	if err != nil {
		t.Fatal(err)
	}

	phKeyHandleClientMac2, phKeyHandleServerMac2, phKeyHandleClientEnc2, phKeyHandleServerEnc2, clientWriteIV2, serverWriteIV2, err := gsdf.SDF_GenerateKeywithECDHE_SSL(phKeyHandle2, pucClientRandom, pucServerRandom, identifier.SGDSM3, 128, 128, 128, 128)
	if err != nil {
		t.Fatal(err)
	}

	msg := make([]byte, 128)
	gsdf.SDF_GenerateRandom(msg)

	clientMac1, err := gsdf.SDF_CalculateMAC(phKeyHandleClientMac1, identifier.SGDSM3, nil, msg)
	if err != nil {
		t.Fatal(err)
	}
	clientMac2, err := gsdf.SDF_CalculateMAC(phKeyHandleClientMac2, identifier.SGDSM3, nil, msg)
	if err != nil {
		t.Fatal(err)
	}

	if bytes.Compare(clientMac1, clientMac2) != 0 {
		t.Fatal("client mac error")
	}

	serverMac1, err := gsdf.SDF_CalculateMAC(phKeyHandleServerMac1, identifier.SGDSM3, nil, msg)
	if err != nil {
		t.Fatal(err)
	}
	serverMac2, err := gsdf.SDF_CalculateMAC(phKeyHandleServerMac2, identifier.SGDSM3, nil, msg)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Compare(serverMac1, serverMac2) != 0 {
		t.Fatal("server mac error")
	}

	encMsg := make([]byte, len(msg))
	decMsg := make([]byte, len(msg))
	err = gsdf.SDF_Encrypt(phKeyHandleClientEnc1, identifier.SGDSM4CBC, clientWriteIV1, msg, &encMsg)
	if err != nil {
		t.Fatal(err)
	}
	err = gsdf.SDF_Decrypt(phKeyHandleClientEnc2, identifier.SGDSM4CBC, clientWriteIV2, encMsg, &decMsg)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Compare(msg, decMsg) != 0 {
		t.Fatal("client enc error")
	}

	err = gsdf.SDF_Encrypt(phKeyHandleServerEnc1, identifier.SGDSM4CBC, serverWriteIV1, msg, &encMsg)
	if err != nil {
		t.Fatal(err)
	}
	err = gsdf.SDF_Decrypt(phKeyHandleServerEnc2, identifier.SGDSM4CBC, serverWriteIV2, encMsg, &decMsg)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Compare(msg, decMsg) != 0 {
		t.Fatal("server enc error")
	}

}