package xcrypto

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"golang.org/x/crypto/pkcs12"
	"io/ioutil"
	"os"
	"strings"
)

const (
	PRIVATE_PEM_BEGIN = "-----BEGIN RSA PRIVATE KEY-----\n"
	PRIVATE_PEM_END   = "\n-----END RSA PRIVATE KEY-----"
	PUBLIC_PEM_BEGIN  = "-----BEGIN PUBLIC KEY-----\n"
	PUBLIC_PEM_END    = "\n-----END PUBLIC KEY-----"
)

type PublicStruct struct {
	publicKey *rsa.PublicKey
}

func (p PublicStruct) Verify(sign, content string, hash crypto.Hash) error {
	return Verify(sign, content, p.publicKey, hash)
}

func (p PublicStruct) Encrypt(content string) (string, error) {
	b, err := rsa.EncryptPKCS1v15(rand.Reader, p.publicKey, []byte(content))
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(b), nil
}

func NewPublicRSA(key string) (p PublicStruct, err error) {
	priKey, err := ParsePublicKey(key)
	if err != nil {
		return
	}
	p.publicKey = priKey
	return
}

type PrivateStruct struct {
	privateKey *rsa.PrivateKey
	PublicStruct
}

func (p PrivateStruct) Sign(content string, hash crypto.Hash) (string, error) {
	return Sign(content, p.privateKey, hash)
}

func (p PrivateStruct) Decrypt(content string) (string, error) {
	bcontent, err := base64.StdEncoding.DecodeString(content)
	if err != nil {
		return "", err
	}
	b, err := rsa.DecryptPKCS1v15(rand.Reader, p.privateKey, bcontent)
	if err != nil {
		return "", err
	}
	return string(b), nil
}

func NewPrivateRSA(key string) (p PrivateStruct, err error) {
	priKey, err := ParsePrivateKey(key)
	if err != nil {
		return
	}
	p.privateKey = priKey
	p.publicKey = &priKey.PublicKey
	return
}

func NewPrivateRSAPfx(key, password string) (p PrivateStruct, err error) {
	priKey, err := ParsePrivateKeyPfx(key, password)
	if err != nil {
		return
	}
	p.privateKey = priKey
	p.publicKey = &priKey.PublicKey
	return
}

func Sign(signContent string, privateKey *rsa.PrivateKey, hash crypto.Hash) (string, error) {
	h := hash.New()
	h.Write([]byte(signContent))
	hashed := h.Sum(nil)

	signature, err := rsa.SignPKCS1v15(rand.Reader, privateKey, hash, hashed)
	if err != nil {
		panic(err)
	}
	return base64.StdEncoding.EncodeToString(signature), nil
}

func Verify(sign, content string, publicKey *rsa.PublicKey, hash crypto.Hash) error {
	h := hash.New()
	h.Write([]byte(content))
	hashed := h.Sum(nil)
	sig, _ := base64.StdEncoding.DecodeString(sign)
	return rsa.VerifyPKCS1v15(publicKey, hash, hashed, sig)
}

func ParsePrivateKey(privateKey string) (*rsa.PrivateKey, error) {
	privateKey = FormatPrivateKey(privateKey)
	// 2、解码私钥字节，生成加密对象
	block, _ := pem.Decode([]byte(privateKey))
	if block == nil {
		return nil, errors.New("private key info error")
	}
	// 3、解析DER编码的私钥，生成私钥对象
	priKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return priKey, nil
}

func ParsePublicKey(publicKey string) (*rsa.PublicKey, error) {
	publicKey = FormatPublicKey(publicKey)
	// 2、解码私钥字节，生成加密对象
	block, _ := pem.Decode([]byte(publicKey))
	if block == nil {
		return nil, errors.New("public key info error")
	}
	// 3、解析DER编码的私钥，生成公钥对象
	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return pubKey.(*rsa.PublicKey), nil
}

func FormatPrivateKey(privateKey string) string {
	if !strings.HasPrefix(privateKey, PRIVATE_PEM_BEGIN) {
		privateKey = PRIVATE_PEM_BEGIN + privateKey
	}
	if !strings.HasSuffix(privateKey, PRIVATE_PEM_END) {
		privateKey = privateKey + PRIVATE_PEM_END
	}
	return privateKey
}

func FormatPublicKey(publicKey string) string {
	if !strings.HasPrefix(publicKey, PUBLIC_PEM_BEGIN) {
		publicKey = PUBLIC_PEM_BEGIN + publicKey
	}
	if !strings.HasSuffix(publicKey, PUBLIC_PEM_END) {
		publicKey = publicKey + PUBLIC_PEM_END
	}
	return publicKey
}

func ParsePrivateKeyPfx(privateKeyName, privatePassword string) (*rsa.PrivateKey, error) {
	f, err := os.Open(privateKeyName)
	if err != nil {
		return nil, err
	}

	bytes, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	// 因为pfx证书公钥和密钥是成对的，所以要先转成pem.Block
	blocks, err := pkcs12.ToPEM(bytes, privatePassword)
	if err != nil {
		return nil, err
	}
	//if len(blocks) != 2 {
	//	return nil, errors.New("解密错误")
	//}
	// 拿到第一个block，用x509解析出私钥（当然公钥也是可以的）
	privateKey, err := x509.ParsePKCS1PrivateKey(blocks[0].Bytes)
	if err != nil {
		return nil, err
	}
	return privateKey, nil
}