package crypto

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/pem"
	"os"
)

const size = 1024

func (_rsa *CryptoUtil) GenerateRSAKey() ([]byte, []byte, error) {
	priv, err := rsa.GenerateKey(rand.Reader, size)
	if err != nil {
		return nil, nil, err
	}
	privKey, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return nil, nil, err
	}
	pubKey, err := x509.MarshalPKIXPublicKey(&priv.PublicKey)
	if err != nil {
		return nil, nil, err
	}
	return privKey, pubKey, nil
}

func (_rsa *CryptoUtil) GenerateRSAKeyFile(privateKeyFile string, publicKeyFile string) error {
	privKey, pubKey, err := _rsa.GenerateRSAKey()
	if err != nil {
		return err
	}
	blockPriv := pem.Block{
		Type:    "RSA PKCS8 PRIVATE KEY",
		Headers: nil,
		Bytes:   pubKey,
	}
	blockPub := pem.Block{
		Type:    "RSA PKIX PUBLIC KEY",
		Headers: nil,
		Bytes:   privKey,
	}

	privFile, err := os.Create(privateKeyFile)
	if err != nil {
		return err
	}
	pubFile, err := os.Create(publicKeyFile)
	if err != nil {
		return err
	}
	err = pem.Encode(privFile, &blockPriv)

	if err != nil {
		return err
	}
	err = pem.Encode(pubFile, &blockPub)
	if err != nil {
		return err
	}
	defer func(privFile *os.File) {
		_ = privFile.Close()
	}(privFile)
	defer func(pubFile *os.File) {
		_ = pubFile.Close()
	}(pubFile)
	return nil
}

func (_rsa *CryptoUtil) GenerateRSAKeyHex() (string, string, error) {
	privKey, pubKey, err := _rsa.GenerateRSAKey()
	if err != nil {
		return "", "", err
	}
	return hex.EncodeToString(privKey), hex.EncodeToString(pubKey), nil
}

func (_rsa *CryptoUtil) GenerateRSAKeyBase64() (string, string, error) {
	privKey, pubKey, err := _rsa.GenerateRSAKey()
	if err != nil {
		return "", "", err
	}
	return base64.StdEncoding.EncodeToString(privKey), base64.StdEncoding.EncodeToString(pubKey), nil
}

func (_rsa *CryptoUtil) ParseRSAPrivateKeyStr(privateKey string) (*rsa.PrivateKey, error) {
	var der []byte
	var err error
	if isHexStr(privateKey) {
		der, err = hex.DecodeString(privateKey)
		if err != nil {
			return nil, err
		}
	} else {
		der, err = base64.StdEncoding.DecodeString(privateKey)
		if err != nil {
			return nil, err
		}
	}
	privKey, err := x509.ParsePKCS8PrivateKey(der)
	if err != nil {
		return nil, err
	}
	return privKey.(*rsa.PrivateKey), nil

}

func (_rsa *CryptoUtil) ParseRSAPublicKeyStr(publicKey string) (*rsa.PublicKey, error) {
	var der []byte
	var err error
	if isHexStr(publicKey) {
		der, err = hex.DecodeString(publicKey)
		if err != nil {
			return nil, err
		}
	} else {
		der, err = base64.StdEncoding.DecodeString(publicKey)
		if err != nil {
			return nil, err
		}
	}
	pub, err := x509.ParsePKIXPublicKey(der)
	if err != nil {
		return nil, err
	}
	return pub.(*rsa.PublicKey), nil
}

func (_rsa *CryptoUtil) EncryptRSA(pub *rsa.PublicKey, msg []byte) ([]byte, error) {
	res, err := rsa.EncryptPKCS1v15(rand.Reader, pub, msg)
	return res, err
}

func (_rsa *CryptoUtil) EncryptRSAStr(pub *rsa.PublicKey, msg string) ([]byte, error) {
	return _rsa.EncryptRSA(pub, []byte(msg))
}

func (_rsa *CryptoUtil) EncryptRSAHex(pub *rsa.PublicKey, msg []byte) (string, error) {
	res, err := _rsa.EncryptRSA(pub, msg)
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(res), nil
}

func (_rsa *CryptoUtil) EncryptRSAHexStr(pub *rsa.PublicKey, msg string) (string, error) {
	res, err := _rsa.EncryptRSA(pub, []byte(msg))
	if err != nil {
		return "", err
	}
	return hex.EncodeToString(res), nil
}

func (_rsa *CryptoUtil) EncryptRSABase64(pub *rsa.PublicKey, msg []byte) (string, error) {
	res, err := _rsa.EncryptRSA(pub, msg)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(res), nil
}

func (_rsa *CryptoUtil) EncryptRSABase64Str(pub *rsa.PublicKey, msg string) (string, error) {
	res, err := _rsa.EncryptRSA(pub, []byte(msg))
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(res), nil
}

func (_rsa *CryptoUtil) DecryptRSA(priv *rsa.PrivateKey, ciphertext []byte) ([]byte, error) {
	res, err := rsa.DecryptPKCS1v15(rand.Reader, priv, ciphertext)
	if err != nil {
		return nil, err
	}
	return res, nil
}

func (_rsa *CryptoUtil) DecryptRSAStr(priv *rsa.PrivateKey, ciphertext string) (string, error) {
	if isHexStr(ciphertext) {
		return _rsa.DecryptRSAHexStr(priv, ciphertext)
	} else {
		return _rsa.DecryptRSABase64Str(priv, ciphertext)
	}
}
func (_rsa *CryptoUtil) DecryptRSABase64Str(priv *rsa.PrivateKey, ciphertext string) (string, error) {
	dec, err := base64.StdEncoding.DecodeString(ciphertext)
	if err != nil {
		return "", err
	}
	res, err := rsa.DecryptPKCS1v15(rand.Reader, priv, dec)
	if err != nil {
		return "", err
	}
	return string(res), nil
}
func (_rsa *CryptoUtil) DecryptRSAHexStr(priv *rsa.PrivateKey, ciphertext string) (string, error) {
	dec, err := hex.DecodeString(ciphertext)
	if err != nil {
		return "", err
	}
	res, err := rsa.DecryptPKCS1v15(rand.Reader, priv, dec)
	if err != nil {
		return "", err
	}
	return string(res), nil
}