package crypto

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

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/pkg/errors"
)

func EncodePublicKeyToPEM(key crypto.PublicKey) ([]byte, error) {
	pub, err := x509.MarshalPKIXPublicKey(key)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	data := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pub,
	})
	return data, nil
}

func DecodePEMToPublicKey(pem []byte) (crypto.PublicKey, error) {
	return jwt.ParseRSAPublicKeyFromPEM(pem)
}

func DecodePEMEncryptedPrivateKey(key []byte, passphrase []byte) (*rsa.PrivateKey, error) {
	// Parse PEM block
	var block *pem.Block
	if block, _ = pem.Decode(key); block == nil {
		return nil, errors.New("invalid PEM block")
	}

	decryptedBlock, err := x509.DecryptPEMBlock(block, passphrase)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	var parsedKey interface{}
	if parsedKey, err = x509.ParsePKCS1PrivateKey(decryptedBlock); err != nil {
		return nil, errors.WithStack(err)
	}

	var privateKey *rsa.PrivateKey
	var ok bool
	if privateKey, ok = parsedKey.(*rsa.PrivateKey); !ok {
		return nil, errors.New("invalid RSA private key")
	}

	return privateKey, nil
}

func EncodePrivateKeyToEncryptedPEM(key *rsa.PrivateKey, passphrase []byte) ([]byte, error) {
	if passphrase == nil {
		return nil, errors.New("passphrase cannot be empty")
	}

	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}

	block, err := x509.EncryptPEMBlock(
		rand.Reader, block.Type,
		block.Bytes, passphrase, x509.PEMCipherAES256,
	)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return pem.EncodeToMemory(block), nil
}