package rsa

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

// RSA生成公私密钥
func RSAGenKey(bits int) (pub string, pri string, ok bool) {
	if bits%1024 != 0 {
		return
	}
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return
	}
	privateStream := x509.MarshalPKCS1PrivateKey(privateKey)
	block1 := pem.Block{Type: "private key", Bytes: privateStream}
	pri = string(pem.EncodeToMemory(&block1))
	publicKey := privateKey.PublicKey
	publicStream, err := x509.MarshalPKIXPublicKey(&publicKey)
	if err != nil {
		return
	}
	block2 := pem.Block{Type: "public key", Bytes: publicStream}
	pub = string(pem.EncodeToMemory(&block2))
	ok = true
	return
}

// RSA加密
func Encrypt(pubkey, data string) string {
	block, _ := pem.Decode([]byte(pubkey))
	if block == nil {
		return ""
	}
	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return ""
	}
	pub := pubInterface.(*rsa.PublicKey)
	res, err := rsa.EncryptPKCS1v15(rand.Reader, pub, []byte(data))
	if err != nil {
		return ""
	}
	return base64.StdEncoding.EncodeToString(res)
}
func EncryptWithPublicKey(publicKeyPEM string, plaintext string) (string, error) {
	// 解码PEM格式的公钥
	// block, _ := pem.Decode([]byte(publicKeyPEM))
	// if block == nil {
	// 	return "", nil
	// }
	pubKey := []byte(publicKeyPEM)
	// 解析公钥
	pub, err := x509.ParsePKIXPublicKey(pubKey)
	if err != nil {
		return "", err
	}
	rsaPub, ok := pub.(*rsa.PublicKey)
	if !ok {
		return "", errors.New("not a valid RSA public key")
	}
	// 将字符串转换为字节数组
	message := []byte(plaintext)
	// 使用OAEP填充和SHA-256哈希函数进行加密
	label := []byte("") // OAEP label, 可以根据需要设置
	hash := sha256.New()
	ciphertext, err := rsa.EncryptOAEP(hash, rand.Reader, rsaPub, message, label)
	if err != nil {
		return "", err
	}
	// 将密文编码为base64字符串
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// RSA解密
func Decrypt(prikey, data string) string {
	if len(data) < 4 {
		return ""
	}
	ciphertext, err := base64.StdEncoding.DecodeString(data)
	if err != nil {
		return ""
	}
	block, _ := pem.Decode([]byte(prikey))
	if block == nil {
		return ""
	}
	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return ""
	}
	text, err := rsa.DecryptPKCS1v15(rand.Reader, priv, ciphertext)
	if err != nil {
		return ""
	}
	return string(text)
}