package crypto

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"fmt"
)

// RSAEncryptor RSA加密器
type RSAEncryptor struct {
	publicKey  *rsa.PublicKey
	privateKey *rsa.PrivateKey
}

// NewRSAEncryptor 创建RSA加密器
func NewRSAEncryptor(publicKeyPEM, privateKeyPEM []byte) (*RSAEncryptor, error) {
	encryptor := &RSAEncryptor{}

	if len(publicKeyPEM) > 0 {
		block, _ := pem.Decode(publicKeyPEM)
		if block == nil {
			return nil, fmt.Errorf("解析公钥失败")
		}

		pub, err := x509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("解析公钥失败: %w", err)
		}

		var ok bool
		encryptor.publicKey, ok = pub.(*rsa.PublicKey)
		if !ok {
			return nil, fmt.Errorf("不是RSA公钥")
		}
	} else {
		return nil, fmt.Errorf("公钥未设置")
	}

	if len(privateKeyPEM) > 0 {
		block, _ := pem.Decode(privateKeyPEM)
		if block == nil {
			return nil, fmt.Errorf("解析私钥失败")
		}

		priv, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("解析私钥失败: %w", err)
		}

		var ok bool
		encryptor.privateKey, ok = priv.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("不是RSA私钥")
		}
	} else {
		return nil, fmt.Errorf("私钥未设置")
	}

	return encryptor, nil
}

func (e *RSAEncryptor) Encrypt(plaintext []byte) ([]byte, error) {
	if e.publicKey == nil {
		return nil, fmt.Errorf("公钥未设置")
	}

	return rsa.EncryptOAEP(sha256.New(), rand.Reader, e.publicKey, plaintext, nil)
}

func (e *RSAEncryptor) Decrypt(ciphertext []byte) ([]byte, error) {
	if e.privateKey == nil {
		return nil, fmt.Errorf("私钥未设置")
	}

	return rsa.DecryptOAEP(sha256.New(), rand.Reader, e.privateKey, ciphertext, nil)
}

func (e *RSAEncryptor) Name() string {
	return "RSA-OAEP-SHA256"
}