package utils

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
)

// LoadCertificate 通过证书的文本内容加载证书
func LoadCertificate(certificateStr string) (certificate *x509.Certificate, err error) {
	block, _ := pem.Decode([]byte(certificateStr))
	if block == nil {
		return nil, fmt.Errorf("decode certificate error: %v", err)
	}
	if block.Type != "CERTIFICATE" {
		return nil, fmt.Errorf("the kind of PEM should be CERTIFICATE")
	}
	certificate, err = x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parse certificate error: %s", err.Error())
	}
	return certificate, nil
}

// LoadPrivateKey 通过私钥的文本内容加载私钥
func LoadPrivateKey(privateKeyStr string) (privateKey *rsa.PrivateKey, err error) {
	block, _ := pem.Decode([]byte(privateKeyStr))
	if block == nil {
		return nil, fmt.Errorf("decode private key err")
	}
	if block.Type != "PRIVATE KEY" {
		return nil, fmt.Errorf("the kind of PEM should be PRVATE KEY")
	}
	key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parse private key err:%s", err.Error())
	}
	privateKey, ok := key.(*rsa.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("not a RSA private key")
	}
	return privateKey, nil
}

// LoadPublicKey 通过公钥文本内容加载公钥
func LoadPublicKey(publicKeyStr string) (publicKey *rsa.PublicKey, err error) {
	block, _ := pem.Decode([]byte(publicKeyStr))
	if block == nil {
		return nil, errors.New("decode public key error")
	}
	if block.Type != "PUBLIC KEY" {
		return nil, fmt.Errorf("the kind of PEM should be PUBLIC KEY")
	}
	key, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parse public key error: %v", err.Error())
	}
	publicKey, ok := key.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("%s is not rsa public key", publicKeyStr)
	}
	return publicKey, nil
}

// GenCertificate 生成bits大小的公钥和私钥文本内容(推荐3096大小)
func GenCertificate(bits int) (error error, privateKeyStr, publicKeyStr string) {
	// 生成私钥文件
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return err, "", ""
	}

	privateKeyDerStream, err := x509.MarshalPKCS8PrivateKey(privateKey)
	privateKeyBlock := &pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: privateKeyDerStream,
	}

	var privateKeyBytes bytes.Buffer
	err = pem.Encode(&privateKeyBytes, privateKeyBlock)
	if err != nil {
		return err, "", ""
	}

	// 生成公钥文件
	publicKey := &privateKey.PublicKey
	publicKeyDerStream, err := x509.MarshalPKIXPublicKey(publicKey)
	publicBlock := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: publicKeyDerStream,
	}

	var publicKeyBytes bytes.Buffer
	err = pem.Encode(&publicKeyBytes, publicBlock)
	if err != nil {
		return err, "", ""
	}

	return nil, privateKeyBytes.String(), publicKeyBytes.String()
}

// GenCertificateFile 生成证书文件
func GenCertificateFile(bits int) (err error) {
	// 生成rsa秘钥
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return err
	}

	// 生成私钥
	privateKeyDerStream, err := x509.MarshalPKCS8PrivateKey(privateKey)
	privateKeyBlock := &pem.Block{
		Type:  "PRIVATE KEY",
		Bytes: privateKeyDerStream,
	}

	privateFile, err := os.Create("private.pem")
	if err != nil {
		return err
	}

	err = pem.Encode(privateFile, privateKeyBlock)
	if err != nil {
		return err
	}

	// 生成公钥文件
	publicKey := &privateKey.PublicKey
	publicKeyDerStream, err := x509.MarshalPKIXPublicKey(publicKey)
	publicBlock := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: publicKeyDerStream,
	}

	publicFile, err := os.Create("public.pem")
	if err != nil {
		return err
	}

	err = pem.Encode(publicFile, publicBlock)
	if err != nil {
		return err
	}

	return nil

}

// LoadPrivateKeyWithPath 通过私钥的文件加载私钥
func LoadPrivateKeyWithPath(fileName string) (privateKey *rsa.PrivateKey, err error) {
	privateKeyBytes, err := ioutil.ReadFile(fileName)
	if err != nil {
		return nil, fmt.Errorf("read private pem file err: %s", err)
	}
	return LoadPrivateKey(string(privateKeyBytes))
}

// LoadPublicKeyWithPath 通过公钥文件加载公钥
func LoadPublicKeyWithPath(fileName string) (privateKey *rsa.PublicKey, err error) {
	publicKeyBytes, err := ioutil.ReadFile(fileName)
	if err != nil {
		return nil, fmt.Errorf("read public pem file err: %s", err)
	}
	return LoadPublicKey(string(publicKeyBytes))
}