package rsa

import (
	"bytes"
	"crypto"
	cryptoRSA "crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"testing"
	"time"

	"xdx.jelly/xgcl/api/common"
)

func BenchmarkKey(t *testing.B) {
	_, _, _ = PKCS1v15{}.GenerateKeyPairRSA(2048)
}
func TestKeyGen(t *testing.T) {

	sk, pk, _ := PKCS1v15{}.GenerateKeyPairRSA(1024)
	fmt.Println(sk.D.Text(16))
	fmt.Println(pk.E)
	fmt.Println(pk.N.Text(16))
}

func TestKey(t *testing.T) {
	sk, pk, _ := PKCS1v15{}.GenerateKeyPairRSA(2048)
	sk1 := sk
	pk1 := pk
	var sdfPriv common.RSArefPrivateKey
	var sdfPub common.RSArefPublicKey
	if err := sk.MarshalSDF(&sdfPriv); err != nil {
		t.Log("sk.MarshalSDF failed")
		t.Fail()
	}
	if err := sk.UnmarshalSDF(&sdfPriv); err != nil || sk != sk1 {
		t.Log("sk.UnmarshalSDF failed")
		t.Fail()
	}

	if err := pk.MarshalSDF(&sdfPub); err != nil {
		t.Log("pk.MarshalSDF failed")
		t.Fail()
	}
	if err := pk.UnmarshalSDF(&sdfPub); err != nil || pk != pk1 {
		t.Log("pk.UnmarshalSDF failed")
		t.Fail()
	}

}

//RAW签摘要
func TestRSARAWSign(t *testing.T) {
	// sha256摘要
	msg := []byte("message")
	digestData := sha256.Sum256(msg)

	sk, pk, err := PKCS1v15{}.GenerateKeyPairRSA(2048)
	fmt.Printf("N:%s\n", sk.N.Text(16)) // sig is the same each time

	// 签名值。2048bit integer
	signData, err := PKCS1v15{}.Sign(sk, crypto.SHA256, digestData[:])
	fmt.Printf("sig:%02x\n", signData) // sig is the same each time
	err = PKCS1v15{}.Verify(pk, crypto.SHA256, digestData[:], signData)
	if err != nil {
		t.Log("verify failed")
		t.Fail()
	}

}

func rsaRAWVerifyDigest(certPEM string, signData []byte, digestData []byte) error {
	block, _ := pem.Decode([]byte(certPEM))
	if block == nil {
		return fmt.Errorf("failed to decode certificate PEM")
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return fmt.Errorf("failed to parse certificate PEM")
	}

	if cert.PublicKeyAlgorithm != x509.RSA {
		return fmt.Errorf("not rsa algorithm")
	}

	pk, ok := cert.PublicKey.(*cryptoRSA.PublicKey)
	if !ok {
		return fmt.Errorf("failed to get publickey")
	}

	var hash crypto.Hash
	switch cert.SignatureAlgorithm {
	case x509.SHA1WithRSA:
		hash = crypto.SHA1
	case x509.SHA256WithRSA:
		hash = crypto.SHA256
	case x509.SHA384WithRSA:
		hash = crypto.SHA384
	case x509.SHA512WithRSA:
		hash = crypto.SHA512
	}

	return PKCS1v15{}.Verify((&PublicKey{}).UnmarshalCryptoRsa(pk), hash, digestData, signData)
}

//RAW验摘要
func TestRSARAWVerifyDigest(t *testing.T) {
	digestData, _ := base64.StdEncoding.DecodeString("ungWv48Bz+pBQUDeXa4iI7ADYaOWF3qctBD/YfIAFa0=")
	signData, _ := base64.StdEncoding.DecodeString("OkXl4owB9VncVmnUXSgQxwMwy8KVuFXk1e9vtOdrXEAOiAXbc6WYjmbT4tdoiEIj02ZaAT449QxQ2fkAsD4IP6iyrmP4Q34qlD5mbh3unuAchLW+XYuDFzVLMp2BLBw2aN/OmjaQdZo6GKy6h5RvrXgf+FwRwI1RIi4r+heA3mlnkUejSJXMkXsbamxMfTO+DH9GBv5cbSwLOEyGcnq/M7urIo/esZT5f+RXwIiI1Z+Qs+1UQ8WJ5O9oBgkJp4BY+o9oaEPuEIMsR0Hz94zv+9+CGn1G0xoBYVOxlRLWYsCyRYnwXzL/2jkJ3GAY3HQyiGa8WRB2LT+C2ijNK3D/nA==")
	certPEM := "MIIDiTCCAnGgAwIBAgIKb0XKECkuAA7e5jANBgkqhkiG9w0BAQsFADAlMRUwEwYDVQQDDAxBQkMgVEVTVCBDQTIxDDAKBgNVBAoMA0FCQzAeFw0xODA0MDIwMDQ4MzFaFw0yMzA0MDIwMDQ4MzFaME4xLzAtBgNVBAMMJjQ0OTk0NDIxMzAwLjAwMDUuMDAwMC42NjQ0NTAyNTY2MjY5NTAxMQ0wCwYDVQQLDARVd2luMQwwCgYDVQQKDANBQkMwggEiMA0GCSqGSIb3DQEBAQUAA4IBDwAwggEKAoIBAQCHbgJr3CmQf69XOK27F4V3YDCfvnH2JcL4xDxHd2B2Sy29NeU4bFsCZaBTc2N3ujHTLSRn+NEuFE90aejd/TvvuDzKda0r0615wo5NUfQiY5pi7UQ3k1IXtSpuImlIdPhnS0egBHropMQgd1f1K+54bSWB9+vhCIATP9GDQS32fJ2z3If0E3T5fCKiS4KS+SBhwtvEEMisCbDpD2y+BwYz0Ozkth9SrJ6r+ujLGwu5JYD/YRL7or8lpi575KvT6Ahf0iDZ/bzLjyZ5rWjSAWWAPJCjUoOFDt9XhY5fTo4B0lj6RV96qyCIau8n+fSVIiqXfxK5/W+iHj9zTXfe2sNbAgMBAAGjgZEwgY4wHwYDVR0jBBgwFoAUErtyLm5Iqy6z01KtzsJ33EmoKtIwCQYDVR0TBAIwADBBBgNVHR8EOjA4MDagNKAypDAwLjEQMA4GA1UEAwwHY3JsNDQ5MjEMMAoGA1UECwwDY3JsMQwwCgYDVQQKDANBQkMwHQYDVR0OBBYEFAHfn4lC2TqykHQ9dIhle8y1d7TyMA0GCSqGSIb3DQEBCwUAA4IBAQAwjLw04jTEDh06UepD7pjmH7rcmqOeXMsfPRLPF1vG2w8dOAvwH7fUIRYWKx4i9uM/mue5rYZokP5p8bV/Gx1+KNEEK7LOgUKlHjc5U/HVa89SnynWCUs9tD1DDM25BSb9szzwUz0QG5XNgskn7h3LsXoUdzXV+j+eIYuIY3q/TURN/3ds7if+0zwnaFwj9PpUcB3A9ShynhpZX6xBLMH3WT8MofUdzVmvZDPxaGxBpORQL8ZlXdU34aeHX9290lBTqH3TxXXEuWoBfzW2cp0j5daoGh9VX7Qh93d3mj+vZKv/W6B0GmNfEiABCNUrN1x4nz7jK1yyMu1j9rJTSXOQ"
	certPEM = "-----BEGIN CERTIFICATE-----\n" + certPEM + "\n-----END CERTIFICATE-----"
	err := rsaRAWVerifyDigest(certPEM, signData, digestData)
	if err != nil {
		t.Log("verify failed")
		t.Fail()
	}
}

func TestSpeed(t *testing.T) {
	msg := []byte("message")
	digestData := sha256.Sum256(msg)

	sk, pk, _ := PKCS1v15{}.GenerateKeyPairRSA(2048)
	// sk.Precompute()
	var signData []byte

	for i := 0; i < 1000; i++ {
		signData, _ = PKCS1v15{}.Sign(sk, crypto.SHA256, digestData[:])
	}

	cnt := 1000
	start := time.Now()
	for i := 0; i < cnt; i++ {
		signData, _ = PKCS1v15{}.Sign(sk, crypto.SHA256, digestData[:])
	}
	end := time.Now()
	elapsed := end.Sub(start)
	fmt.Printf("Sign %d, used time: %d ms, %d pcs/s\n", cnt, elapsed.Milliseconds(), int(float64(cnt)/float64(elapsed.Milliseconds())*1000))

	start = time.Now()
	for i := 0; i < cnt; i++ {
		_ = PKCS1v15{}.Verify(pk, crypto.SHA256, digestData[:], signData)
	}
	end = time.Now()
	elapsed = end.Sub(start)
	fmt.Printf("verify %d, used time: %d ms, %d pcs/s\n", cnt, elapsed.Milliseconds(), int(float64(cnt)/float64(elapsed.Milliseconds())*1000))

}

func BenchmarkSign(b *testing.B) {
	msg := []byte("message")
	digestData := sha256.Sum256(msg)

	sk, _, _ := PKCS1v15{}.GenerateKeyPairRSA(1024 * 2)
	b.StartTimer()
	_, _ = PKCS1v15{}.Sign(sk, crypto.SHA256, digestData[:])
	b.StopTimer()
}

func BenchmarkVerify(b *testing.B) {
	msg := []byte("message")
	digestData := sha256.Sum256(msg)

	sk, pk, _ := PKCS1v15{}.GenerateKeyPairRSA(2048)
	signData, _ := PKCS1v15{}.Sign(sk, crypto.SHA256, digestData[:])
	b.StartTimer()

	_ = PKCS1v15{}.Verify(pk, crypto.SHA256, digestData[:], signData)
	b.StopTimer()
}

func TestRsaSdf(t *testing.T) {
	impl := PKCS1v15{}

	sk, pk, err := impl.GenerateKeyPairRSA(2048)
	if err != nil || sk.Validate() != nil {
		t.Log("GenerateKeyPairRSA failed")
		t.Fail()
		return
	}

	envelopedKeyPair, err := impl.GenerateEnvelopedKeyPairRSA(1024, pk)
	if err != nil {
		t.Log("GenerateEnvelopedKeyPairRSA failed")
		t.Fail()
		return
	}

	sk1, pk1, err := impl.ImportEnvelopedKeyPairRSA(envelopedKeyPair, sk)
	if err != nil || sk1.Validate() != nil {
		t.Log("ImportEnvelopedKeyPairRSA failed", err)
		t.Fail()
		return
	}

	sessionKey, encryptedSessionKey, err := impl.GenerateKeyRSA(16, pk)
	if err != nil {
		t.Log("GenerateKeyRSA failed")
		t.Fail()
		return
	}

	encryptedSessionKey, err = impl.ExchangeDigitEnvelopeRSA(sk, pk1, encryptedSessionKey)
	if err != nil {
		t.Log("ExchangeDigitEnvelopeRSA failed")
		t.Fail()
		return
	}

	decrytpedSessionKey, err := impl.ImportKeyRSA(16, sk1, encryptedSessionKey)
	if err != nil {
		t.Log("ImportKeyRSA failed")
		t.Fail()
		return
	}
	if bytes.Compare(sessionKey, decrytpedSessionKey) != 0 {
		t.Log("sessionkey not equal")
		t.Fail()
		return
	}

	data := make([]byte, pk.Size())
	data[0] = 1
	encryptedData, err := impl.PublicKeyOperationRSA(pk, data)
	if err != nil {
		t.Log("PublicKeyOperationRSA failed")
		t.Fail()
		return
	}

	decryptedData, err := impl.PrivateKeyOperationRSA(sk, encryptedData)
	if err != nil {
		t.Log("PrivateKeyOperationRSA failed")
		t.Fail()
		return
	}

	if bytes.Compare(data, decryptedData) != 0 {
		t.Log("PublicKeyOperationRSA/PrivateKeyOperationRSA failed")
		t.Fail()
		return
	}

}