package rsa

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

var (
	ErrPrivateKey = errors.New("private key error")
	ErrPublicKey  = errors.New("public key error")
)

func PrivateKey(pri []byte) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode(pri)
	if block == nil {
		return nil, ErrPrivateKey
	}
	return x509.ParsePKCS1PrivateKey(block.Bytes)
}

func PublicKey(pub []byte) (*rsa.PublicKey, error) {
	block, _ := pem.Decode(pub)
	if block == nil {
		return nil, ErrPublicKey
	}
	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	rsaPub := pubInterface.(*rsa.PublicKey)
	return rsaPub, nil
}

func Encrypt(orig, pubKey []byte) ([]byte, error) {
	pub, err := PublicKey([]byte(pubKey))
	if err != nil {
		return nil, err
	}
	cipher, err := rsa.EncryptPKCS1v15(rand.Reader, pub, orig)
	if err != nil {
		return nil, err
	}

	return cipher, nil
}

func Decrypt(cipher, priKey []byte) ([]byte, error) {
	pri, err := PrivateKey([]byte(priKey))
	if err != nil {
		return nil, err
	}
	ori, err := rsa.DecryptPKCS1v15(nil, pri, cipher)
	if err != nil {
		return nil, err
	}
	return ori, nil
}

func GenerateKey(bits int) (priKey, pubKey []byte, err error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return
	}

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

	publicKey := &privateKey.PublicKey
	derPkix, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		return
	}

	block = &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: derPkix,
	}
	pubKey = pem.EncodeToMemory(block)
	return
}