package utils

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"log"
)

// 生成RSA密钥对
func GenerateRSAKeyPair(bits int) (*rsa.PrivateKey, *rsa.PublicKey) {
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		log.Fatalf("Failed to generate RSA key pair: %v", err)
	}
	return privateKey, &privateKey.PublicKey
}

// 将公钥编码为PEM格式
func EncodePublicKeyToPEM(publicKey *rsa.PublicKey) []byte {
	pubASN1, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		log.Fatalf("Failed to marshal public key: %v", err)
	}
	
	pubPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubASN1,
	})
	
	return pubPEM
}

// 将私钥编码为PEM格式
func EncodePrivateKeyToPEM(privateKey *rsa.PrivateKey) []byte {
	privPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
	})
	
	return privPEM
}

// 从PEM格式解码公钥
func DecodePublicKeyFromPEM(pubPEM []byte) *rsa.PublicKey {
	block, _ := pem.Decode(pubPEM)
	if block == nil {
		log.Fatal("Failed to decode PEM block containing public key")
	}
	
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		log.Fatalf("Failed to parse public key: %v", err)
	}
	
	return pub.(*rsa.PublicKey)
}

// 从PEM格式解码私钥
func DecodePrivateKeyFromPEM(privPEM []byte) *rsa.PrivateKey {
	block, _ := pem.Decode(privPEM)
	if block == nil {
		log.Fatal("Failed to decode PEM block containing private key")
	}
	
	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		log.Fatalf("Failed to parse private key: %v", err)
	}
	
	return priv
}