package rsa

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"io/ioutil"
)

// RSA公钥私钥产生
func GeneratePrivAndPubKey() (prvkey, pubkey []byte) {
	// 生成私钥文件
	privateKey, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		panic(err)
	}
	derStream := x509.MarshalPKCS1PrivateKey(privateKey)
	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: derStream,
	}
	prvkey = pem.EncodeToMemory(block)
	//保存到文件
	err = savePEMToFile(prvkey, "config/private.pem")
	if err != nil {
		panic(err)
	}

	publicKey := &privateKey.PublicKey
	derPkix, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		panic(err)
	}
	block = &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: derPkix,
	}
	pubkey = pem.EncodeToMemory(block)
	err = savePEMToFile(pubkey, "config/public.pem")
	if err != nil {
		panic(err)
	}
	return
}

// 保存文件
func savePEMToFile(data []byte, filename string) error {
	err := ioutil.WriteFile(filename, data, 0644)
	if err != nil {
		return err
	}
	//fmt.Println("Saved", filename)
	return nil
}

// 加密
func Encryption(data []byte) []byte {
	publickKeyDecoded, err2 := GetPubKey()
	if err2 != nil {
		panic(err2)
	}
	encryptedData, err := rsa.EncryptPKCS1v15(rand.Reader, publickKeyDecoded.(*rsa.PublicKey), data)
	if err != nil {
		panic(err)
	}
	return encryptedData
}

// 解密
func Decryption(encryptedData []byte) []byte {
	privateKeyDecoded := GetPrivKey()
	decryptedData, _ := rsa.DecryptPKCS1v15(rand.Reader, privateKeyDecoded, encryptedData)
	return decryptedData
}

// 加载文件
func loadPEMFromFile(filename string) ([]byte, error) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return data, nil
}

// 读取私钥
func GetPrivKey() *rsa.PrivateKey {
	privateKeyBytes, err := loadPEMFromFile("config/private.pem")
	if err != nil {
		panic(err)
	}
	privateKeyBlock, _ := pem.Decode(privateKeyBytes)
	privateKeyDecoded, err := x509.ParsePKCS1PrivateKey(privateKeyBlock.Bytes)
	if err != nil {
		panic(err)
	}
	return privateKeyDecoded
}

// 读取公钥
func GetPubKey() (interface{}, error) {
	publicKeyBytes, err := loadPEMFromFile("config/public.pem")
	if err != nil {
		panic(err)
	}
	publicKeyBlock, _ := pem.Decode(publicKeyBytes)
	publicKeyDecoded, err := x509.ParsePKIXPublicKey(publicKeyBlock.Bytes)

	if err != nil {
		panic(err)
	}
	return publicKeyDecoded, nil
}