package utils

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

// 生成私钥文件 TODO 未指定路径
func RsaKeyGen(bits int) error {
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return err
	}
	derStream := x509.MarshalPKCS1PrivateKey(privateKey)
	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: derStream,
	}
	priFile, err := os.Create("private.pem")
	if err != nil {
		return err
	}
	err = pem.Encode(priFile, block)
	priFile.Close()
	if err != nil {
		return err
	}
	// 生成公钥文件
	publicKey := &privateKey.PublicKey
	derPkix, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		return err
	}
	block = &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: derPkix,
	}
	pubFile, err := os.Create("public.pem")
	if err != nil {
		return err
	}
	err = pem.Encode(pubFile, block)
	pubFile.Close()
	if err != nil {
		return err
	}
	return nil
}

// 生成私钥文件, 返回 privateKey , publicKey, error
func RsaKeyGenText(bits int) (string, string, error) { // bits 字节位 1024/2048
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return "", "", err
	}
	derStream := x509.MarshalPKCS1PrivateKey(privateKey)
	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: derStream,
	}
	priBuff := bytes.NewBuffer(nil)
	err = pem.Encode(priBuff, block)
	if err != nil {
		return "", "", err
	}
	// 生成公钥文件
	publicKey := &privateKey.PublicKey
	derPkix, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		return "", "", err
	}
	block = &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: derPkix,
	}
	pubBuff := bytes.NewBuffer(nil)
	err = pem.Encode(pubBuff, block)
	if err != nil {
		return "", "", err
	}
	return priBuff.String(), pubBuff.String(), nil
}

// 加密
func RsaEncrypt(rawData, publicKey []byte) ([]byte, error) {
	block, _ := pem.Decode(publicKey)
	if block == nil {
		return nil, errors.New("public key error")
	}
	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	pub := pubInterface.(*rsa.PublicKey)
	return rsa.EncryptPKCS1v15(rand.Reader, pub, rawData)
}

// 公钥加密
func RsaEncrypts(data, keyBytes []byte) []byte {
	//解密pem格式的公钥
	block, _ := pem.Decode(keyBytes)
	if block == nil {
		panic(errors.New("public key error"))
	}
	// 解析公钥
	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		panic(err)
	}
	// 类型断言
	pub := pubInterface.(*rsa.PublicKey)
	//加密
	ciphertext, err := rsa.EncryptPKCS1v15(rand.Reader, pub, data)
	if err != nil {
		panic(err)
	}
	return ciphertext
}

// 解密
func RsaDecrypt(cipherText, privateKey []byte) ([]byte, error) {
	block, _ := pem.Decode(privateKey)
	if block == nil {
		return nil, errors.New("private key error")
	}
	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return rsa.DecryptPKCS1v15(rand.Reader, priv, cipherText)
}
func SyRsaEncrypt(signStr, signature, privateKeyBytes string) error {
	h := sha256.New()
	h.Write([]byte(signStr))
	d := h.Sum(nil)
	sign, _ := base64.StdEncoding.DecodeString(signature)
	pub, err := ParsePKIXPublicKey(privateKeyBytes)
	err = rsa.VerifyPKCS1v15(pub, crypto.SHA256, d, sign)
	if err != nil {
		return err
	}

	return nil
}
func FormatPublicKey(publicKey string) (pKey string) {
	var buffer strings.Builder
	buffer.WriteString("-----BEGIN PUBLIC KEY-----\n")
	rawLen := 64
	keyLen := len(publicKey)
	raws := keyLen / rawLen
	temp := keyLen % rawLen
	if temp > 0 {
		raws++
	}
	start := 0
	end := start + rawLen
	for i := 0; i < raws; i++ {
		if i == raws-1 {
			buffer.WriteString(publicKey[start:])
		} else {
			buffer.WriteString(publicKey[start:end])
		}
		buffer.WriteByte('\n')
		start += rawLen
		end = start + rawLen
	}
	buffer.WriteString("-----END PUBLIC KEY-----\n")
	pKey = buffer.String()
	return
}

func ParsePKIXPublicKey(privateKey string) (*rsa.PublicKey, error) {

	privateKey = FormatPublicKey(privateKey)
	// 2、解码私钥字节，生成加密对象
	block, _ := pem.Decode([]byte(privateKey))
	if block == nil {

		return nil, errors.New("私钥信息错误！")
	}
	// 3、解析DER编码的私钥，生成私钥对象
	priKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {

		return nil, err
	}
	return priKey.(*rsa.PublicKey), nil
}

// 从证书获取公钥
func OpensslPemGetPublic(pathOrString string) (interface{}, error) {
	var certPem []byte
	var err error
	if IsFile(pathOrString) && Exists(pathOrString) {
		certPem, err = ioutil.ReadFile(pathOrString)
		if err != nil {
			return nil, err
		}
		if string(certPem) == "" {
			return nil, errors.New("empty pem file")
		}
	} else {
		if pathOrString == "" {
			return nil, errors.New("empty pem string")
		}
		certPem = StringToSlice(pathOrString)
	}
	block, rest := pem.Decode(certPem)
	if block == nil || block.Type != "PUBLIC KEY" {
		//log.Fatal("failed to decode PEM block containing public key")
		return nil, errors.New("failed to decode PEM block containing public key")
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Got a %T, with remaining data: %q", pub, rest)
	return pub, nil
}