package asymm

import (
	"bytes"
	"crypto/rsa"
	"testing"
)

func TestRsaPrivateKeyEncodeDecodeWithLegacyFlag(t *testing.T) {
	priv, pub, err := GenerateRsaKey(1024)
	if err != nil {
		t.Fatalf("GenerateRsaKey failed: %v", err)
	}

	modernPEM, err := EncodeRsaPrivateKeyWithLegacy(priv, "pwd", false)
	if err != nil {
		t.Fatalf("EncodeRsaPrivateKeyWithLegacy modern failed: %v", err)
	}
	if !bytes.Contains(modernPEM, []byte("ENCRYPTED PRIVATE KEY")) {
		t.Fatalf("modern PEM should be ENCRYPTED PRIVATE KEY")
	}
	modernPriv, err := DecodeRsaPrivateKeyWithLegacy(modernPEM, "pwd", false)
	if err != nil {
		t.Fatalf("DecodeRsaPrivateKeyWithLegacy modern failed: %v", err)
	}
	if modernPriv.D.Cmp(priv.D) != 0 {
		t.Fatalf("modern decoded private key mismatch")
	}

	autoPriv, err := DecodeRsaPrivateKey(modernPEM, "pwd")
	if err != nil {
		t.Fatalf("DecodeRsaPrivateKey auto failed: %v", err)
	}
	if autoPriv.D.Cmp(priv.D) != 0 {
		t.Fatalf("auto decoded private key mismatch")
	}

	if _, err := DecodePrivateKey(modernPEM, "pwd"); err != nil {
		t.Fatalf("DecodePrivateKey for encrypted PKCS8 failed: %v", err)
	}

	legacyPEM, err := EncodeRsaPrivateKeyWithLegacy(priv, "pwd", true)
	if err != nil {
		t.Fatalf("EncodeRsaPrivateKeyWithLegacy legacy failed: %v", err)
	}
	if !bytes.Contains(legacyPEM, []byte("RSA PRIVATE KEY")) {
		t.Fatalf("legacy PEM should be RSA PRIVATE KEY")
	}
	legacyPriv, err := DecodeRsaPrivateKeyWithLegacy(legacyPEM, "pwd", true)
	if err != nil {
		t.Fatalf("DecodeRsaPrivateKeyWithLegacy legacy failed: %v", err)
	}
	if legacyPriv.D.Cmp(priv.D) != 0 {
		t.Fatalf("legacy decoded private key mismatch")
	}

	pubPEM, err := EncodeRsaPublicKey(pub)
	if err != nil {
		t.Fatalf("EncodeRsaPublicKey failed: %v", err)
	}
	decodedPub, err := DecodeRsaPublicKey(pubPEM)
	if err != nil {
		t.Fatalf("DecodeRsaPublicKey failed: %v", err)
	}
	if decodedPub.N.Cmp(pub.N) != 0 || decodedPub.E != pub.E {
		t.Fatalf("decoded public key mismatch")
	}
}

func TestRSAOAEPEncryptDecrypt(t *testing.T) {
	priv, pub, err := GenerateRsaKey(1024)
	if err != nil {
		t.Fatalf("GenerateRsaKey failed: %v", err)
	}

	msg := []byte("rsa-oaep-message")
	label := []byte("label")
	enc, err := RSAEncryptOAEP(pub, msg, label, 0)
	if err != nil {
		t.Fatalf("RSAEncryptOAEP failed: %v", err)
	}
	dec, err := RSADecryptOAEP(priv, enc, label, 0)
	if err != nil {
		t.Fatalf("RSADecryptOAEP failed: %v", err)
	}
	if !bytes.Equal(dec, msg) {
		t.Fatalf("oaep decrypt mismatch")
	}
}

func TestRSAPSSSignVerify(t *testing.T) {
	priv, pub, err := GenerateRsaKey(1024)
	if err != nil {
		t.Fatalf("GenerateRsaKey failed: %v", err)
	}

	privPEM, err := EncodeRsaPrivateKeyWithLegacy(priv, "pwd", false)
	if err != nil {
		t.Fatalf("EncodeRsaPrivateKeyWithLegacy failed: %v", err)
	}
	pubPEM, err := EncodeRsaPublicKey(pub)
	if err != nil {
		t.Fatalf("EncodeRsaPublicKey failed: %v", err)
	}

	msg := []byte("rsa-pss-message")
	sig, err := RSASignPSS(msg, privPEM, "pwd", 0, &rsa.PSSOptions{SaltLength: rsa.PSSSaltLengthEqualsHash})
	if err != nil {
		t.Fatalf("RSASignPSS failed: %v", err)
	}
	if err := RSAVerifyPSS(sig, msg, pubPEM, 0, &rsa.PSSOptions{SaltLength: rsa.PSSSaltLengthEqualsHash}); err != nil {
		t.Fatalf("RSAVerifyPSS failed: %v", err)
	}
}

func TestRSAPKCS1v15EncryptDecryptCompatibility(t *testing.T) {
	priv, pub, err := GenerateRsaKey(1024)
	if err != nil {
		t.Fatalf("GenerateRsaKey failed: %v", err)
	}

	msg := []byte("rsa-pkcs1v15-message")
	enc, err := RSAEncrypt(pub, msg)
	if err != nil {
		t.Fatalf("RSAEncrypt failed: %v", err)
	}
	dec, err := RSADecrypt(priv, enc)
	if err != nil {
		t.Fatalf("RSADecrypt failed: %v", err)
	}
	if !bytes.Equal(dec, msg) {
		t.Fatalf("pkcs1v15 decrypt mismatch")
	}
}