package crypt

import (
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"stat/common/log"
)

func splitBytesSlice(data []byte, n int) [][]byte {
	var chunk []byte
	chunks := make([][]byte, 0, len(data)/n+1)
	for len(data) >= n {
		chunk, data = data[:n], data[n:]
		chunks = append(chunks, chunk)
	}
	if len(data) > 0 {
		chunks = append(chunks, data[:])
	}
	return chunks
}

func rsaBytesToPublicKey(pub []byte) (*rsa.PublicKey, error) {
	// fmt.Printf("pubKey is %s", string(pub))
	block, _ := pem.Decode(pub)
	enc := x509.IsEncryptedPEMBlock(block)
	b := block.Bytes
	var err error
	if enc {
		b, err = x509.DecryptPEMBlock(block, nil)
		if err != nil {
			log.Error("RSABytesToPublicKey DecryptPEMBlock err", log.E(err))
			return nil, err
		}
	}

	pk, err := x509.ParsePKIXPublicKey(b)
	if err != nil {
		log.Error("RSABytesToPublicKey ParsePKIXPublicKey err", log.E(err))
		return nil, err
	}
	if pk_, ok := pk.(*rsa.PublicKey); ok {
		return pk_, nil
	}
	log.Error("RSABytesToPublicKey error input not PublicKey")
	return nil, errors.New("Not RSA PublicKey format")
}

func rsaBytesToPrivateKey(priv []byte) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode(priv)
	enc := x509.IsEncryptedPEMBlock(block)
	b := block.Bytes
	var err error
	if enc {
		b, err = x509.DecryptPEMBlock(block, nil)
		if err != nil {
			log.Error("RSABytesToPrivateKey DecryptPEMBlock err", log.E(err))
			return nil, err
		}
	}
	key, err := x509.ParsePKCS8PrivateKey(b)
	if err != nil {
		log.Error("RSABytesToPrivateKey ParsePKCS8PrivateKey err", log.E(err))
		return nil, err
	}
	if pk_, ok := key.(*rsa.PrivateKey); ok {
		return pk_, nil
	}
	log.Error("RSABytesToPrivateKey error input not PrivateKey")
	return nil, errors.New("Not RSA PrivateKey format")
}

// 公钥加密
func RSAEncryptByPublicKey(pemPubKey string, data []byte) ([]byte, error) {
	pubKey, err := rsaBytesToPublicKey([]byte(pemPubKey))
	if err != nil {
		log.Error("RSAEncryptByPublicKey rsaBytesToPublicKey err", log.E(err))
		return nil, err
	}
	partLen := pubKey.N.BitLen()/8 - 11
	chunks := splitBytesSlice(data, partLen)
	result := bytes.Buffer{}
	for _, chunk := range chunks {
		cipherSilce, err := rsa.EncryptPKCS1v15(rand.Reader, pubKey, chunk)
		if err != nil {
			log.Error("RSAEncryptByPublicKey EncryptPKCS1v15 err", log.E(err))
			return nil, err
		}
		result.Write(cipherSilce)
	}
	return result.Bytes(), nil
}

// 私玥签名
func RSASignByPrivateKey(pemPrivateKey string, data []byte, hash crypto.Hash) ([]byte, error) {
	fmt.Printf("pemPrivateKey is %s\n", pemPrivateKey)
	privateKey, err := rsaBytesToPrivateKey([]byte(pemPrivateKey))
	if err != nil {
		log.Error("RSASignByPrivateKey rsaBytesToPrivateKey err", log.E(err))
		return nil, err
	}
	h := hash.New()
	h.Write(data)
	hashed := make([]byte, 0)
	hashed = h.Sum(hashed)

	return rsa.SignPKCS1v15(rand.Reader, privateKey, hash, hashed)

}
func RSAVerifyByPublicKey(publicKey string, data, sign []byte, hash crypto.Hash) (err error) {
	pubKey, err := rsaBytesToPublicKey([]byte(publicKey))
	if err != nil {
		log.Error("RSAEncryptByPublicKey rsaBytesToPublicKey err", log.E(err))
		return err
	}
	h := hash.New()
	h.Write(data)
	hashed := make([]byte, 0)
	hashed = h.Sum(hashed)
	return rsa.VerifyPKCS1v15(pubKey, hash, hashed, sign)
}

// 公钥解密
func RSADecryptByPrivateKey(pemPrivateKey string, cipherData []byte) ([]byte, error) {
	privateKey, err := rsaBytesToPrivateKey([]byte(pemPrivateKey))
	if err != nil {
		log.Error("RSADecryptByPrivateKey rsaBytesToPrivateKey err", log.E(err))
		return nil, err
	}
	blockLen := privateKey.N.BitLen() / 8
	chunks := splitBytesSlice(cipherData, blockLen)
	result := bytes.Buffer{}
	for _, chunk := range chunks {
		plain, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, chunk)
		if err != nil {
			log.Error("RSADecryptByPrivateKey DecryptPKCS1v15 err", log.E(err))
			return nil, err
		}
		result.Write(plain)
	}
	return result.Bytes(), nil
}

// 私玥签名
func RSASignByPrivateKeySHA256(pemPrivateKey string, data []byte, hash crypto.Hash) ([]byte, error) {
	fmt.Printf("pemPrivateKey is %s\n", pemPrivateKey)
	privateKey, err := rsaBytesToPrivateKey([]byte(pemPrivateKey))
	if err != nil {
		log.Error("RSASignByPrivateKey rsaBytesToPrivateKey err", log.E(err))
		return nil, err
	}
	h := sha256.New()
	h.Write(data)
	// hashed := make([]byte, 0)
	hashed := h.Sum(nil)

	return rsa.SignPSS(rand.Reader, privateKey, hash, hashed, nil)
}