package algorithm

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

// PKCS1 for rsa
type PKCS1 struct {
}

func NewPKCS1() *PKCS1 {
	return &PKCS1{}
}

type RsaKey struct {
	Pri io.Writer
	Pub io.Writer
}

// Generate RSA公钥私钥产生
func (*PKCS1) Generate(bits int, out *RsaKey) error {
	// 生成私钥文件
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return err
	}
	derStream := x509.MarshalPKCS1PrivateKey(privateKey)
	priBlock := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: derStream,
	}
	if err = pem.Encode(out.Pri, priBlock); err != nil {
		return err
	}
	// 生成公钥文件
	publicKey := &privateKey.PublicKey
	derPkix, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		return err
	}
	pubBlock := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: derPkix,
	}
	return pem.Encode(out.Pub, pubBlock)
}

// 加密
func (pk *PKCS1) rsaEncrypt(origData []byte, publicKey []byte) ([]byte, error) {
	//解密pem格式的公钥
	block, _ := pem.Decode(publicKey)
	if block == nil {
		return nil, errors.New("public key error")
	}
	// 解析公钥
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	// 类型断言
	pubKey, ok := pub.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("pub case type public key error")
	}
	//加密
	return rsa.EncryptPKCS1v15(rand.Reader, pubKey, origData)
}

// 解密
func (pk *PKCS1) rsaDecrypt(ciphertext []byte, privateKey []byte) ([]byte, error) {
	//解密
	block, _ := pem.Decode(privateKey)
	if block == nil {
		return nil, errors.New("PKCS1: private key error")
	}
	var (
		priv *rsa.PrivateKey
	)
	switch block.Type {
	case "RSA PRIVATE KEY":
		p, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		priv = p

	// RFC5208 - https://tools.ietf.org/html/rfc5208
	case "PRIVATE KEY":
		key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		p, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, errors.New("PKCS1: private key error")
		}
		priv = p

	//case "EC PRIVATE KEY":
	//case "DSA PRIVATE KEY":
	//case "OPENSSH PRIVATE KEY":

	default:
		return nil, fmt.Errorf("unsupported key type %q", block.Type)
	}
	// 解密
	return rsa.DecryptPKCS1v15(rand.Reader, priv, ciphertext)
}

// Encrypt PKCS interface
func (pk *PKCS1) Encrypt(raw, key []byte) (string, error) {
	data, err := pk.rsaEncrypt(raw, key)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(data), nil
}

// Decrypt PKCS interface
func (pk *PKCS1) Decrypt(raw string, key []byte) ([]byte, error) {
	data, err := base64.StdEncoding.DecodeString(raw)
	if err != nil {
		return nil, err
	}
	return pk.rsaDecrypt(data, key)
}

func ReadKey(filename string) ([]byte, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	return io.ReadAll(f)
}