package myth

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

// RSA
type RSA struct {
	publicKey  *rsa.PublicKey
	privateKey *rsa.PrivateKey
}

// NewRSA new rsa
func NewRSA() *RSA {
	return &RSA{}
}

// RSA 公、私钥生成
func RSAGenerateKeys() (privatePEM, publicPEM []byte, err error) {
	var privateKey *rsa.PrivateKey
	privateKey, err = rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return
	}

	// Encode private key to PEM format
	privatePEM = pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
	})

	// Encode public key to PEM format
	var pubASN1 []byte
	pubASN1, err = x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	if err != nil {
		return
	}
	publicPEM = pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: pubASN1,
	})

	return
}

// 载入 公钥
func (r *RSA) LoadPublicKey(data []byte) (err error) {
	var block *pem.Block
	block, _ = pem.Decode(data)
	if block == nil {
		err = errors.New("decode public key error")
		return
	}
	var k any
	k, err = x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return
	}
	r.publicKey = k.(*rsa.PublicKey)
	return
}

// 载入 私钥
func (r *RSA) LoadPrivateKey(data []byte) (err error) {
	var block *pem.Block
	block, _ = pem.Decode(data)
	if block == nil {
		err = errors.New("decode private key error")
		return
	}

	r.privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
	return
}

/*
// Encrypt encrypts data using rsa public key.
func (r *RSA) Encrypt(msg, label []byte) ([]byte, error) {
	return rsa.EncryptOAEP(sha256.New(), rand.Reader, r.publicKey, msg, label)
}

// Encrypt encrypts data using rsa public key.
func (r *RSA) Decrypt(ciphertext, label []byte) ([]byte, error) {
	return rsa.DecryptOAEP(sha256.New(), rand.Reader, r.privateKey, ciphertext, label)
}
	// */

// Sign signs using private key with sha256
func (r *RSA) Sign(data []byte) (signs []byte, err error) {
	h := sha256.Sum256(data)
	signs, err = rsa.SignPKCS1v15(rand.Reader, r.privateKey, crypto.SHA256, h[:])
	return
}

// Sign256 signs using private key with sha512
func (r *RSA) Sign512(data []byte) (signs []byte, err error) {
	h := sha512.Sum512(data)
	signs, err = rsa.SignPKCS1v15(rand.Reader, r.privateKey, crypto.SHA512, h[:])
	return
}

// verifies signature using public key with sha256
func (r *RSA) Verify256(data, signs []byte) (err error) {
	h := sha256.Sum256(data)
	err = rsa.VerifyPKCS1v15(r.publicKey, crypto.SHA256, h[:], signs)
	return
}

// verifies signature using public key with sha512
func (r *RSA) Verify512(data, signs []byte) (err error) {
	h := sha256.Sum256(data)
	err = rsa.VerifyPKCS1v15(r.publicKey, crypto.SHA256, h[:], signs)
	return
}