package encryptUtil

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/md5"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"gitee.com/ccait-dev/fast-api/fast/utils/loggerUtil"
	"io"
	"log"
	"os"
)

var (
	AESIVKey = []byte("12345678901234567890123456789012") // 256位密钥
)

func Sha2(txt string) string {
	hash := sha256.Sum256([]byte(txt))
	return fmt.Sprintf("%x\n", hash)
}

func Md5(txt string) string {
	hash := md5.New()
	io.WriteString(hash, txt)
	md5Sum := hash.Sum(nil)
	txt = fmt.Sprintf("%X", md5Sum)
	return txt
}

func EncryptRSA(txt string) ([]byte, *rsa.PublicKey, *rsa.PrivateKey) {
	// 生成RSA密钥对
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		log.Println(err)
		return nil, nil, nil
	}
	publicKey := &privateKey.PublicKey

	// 对原始数据进行加密
	message := []byte(txt)
	encrypted, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, message)
	if err != nil {
		log.Println(err)
		return nil, nil, nil
	}

	// 输出加密结果
	fmt.Printf("Encrypted: %x\n", encrypted)

	return encrypted, publicKey, privateKey
}

func DecryptRSA(encrypted []byte, privateKey *rsa.PrivateKey) []byte {
	// 使用私钥解密数据
	decrypted, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, encrypted)
	if err != nil {
		loggerUtil.Error(err.Error())
		return nil
	}

	// 输出解密结果
	fmt.Printf("Decrypted: %s\n", decrypted)
	return decrypted
}

func SaveRSAPrivateKey(private *rsa.PrivateKey) {
	out := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(private),
	}
	file, err := os.Create("private.pem")
	if err != nil {
		loggerUtil.Error(err.Error())
		return
	}
	defer file.Close()
	err = pem.Encode(file, out)
	if err != nil {
		loggerUtil.Error(err.Error())
		return
	}
}

func SaveRSAPublicKey(public *rsa.PublicKey) {
	publicBytes, err := x509.MarshalPKIXPublicKey(public)
	if err != nil {
		loggerUtil.Error(err.Error())
		return
	}
	out := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: publicBytes,
	}
	file, err := os.Create("public.pem")
	if err != nil {
		loggerUtil.Error(err.Error())
		return
	}
	defer file.Close()
	err = pem.Encode(file, out)
	if err != nil {
		loggerUtil.Error(err.Error())
		return
	}
}

func EncryptAES(key []byte, plaintext string) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	if key == nil || len(key) == 0 {
		key = AESIVKey
	}
	plainBytes := []byte(plaintext)
	// 对于CBC模式，需要使用PKCS#7填充plaintext到blocksize的整数倍
	plainBytes = pad(plainBytes, aes.BlockSize)
	ciphertext := make([]byte, aes.BlockSize+len(plainBytes))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", err
	}
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext[aes.BlockSize:], plainBytes)
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

func DecryptAES(key []byte, ct string) (string, error) {
	ciphertext, err := base64.StdEncoding.DecodeString(ct)
	if err != nil {
		return "", err
	}
	if key == nil || len(key) == 0 {
		key = AESIVKey
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}
	if len(ciphertext) < aes.BlockSize {
		return "", err
	}
	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(ciphertext, ciphertext)
	// 删除PKCS#7填充
	plaintext := unpad(ciphertext)
	return string(plaintext), nil
}

// pad 使用PKCS#7标准填充数据
func pad(buf []byte, blockSize int) []byte {
	padding := blockSize - (len(buf) % blockSize)
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(buf, padtext...)
}

// unpad 删除PKCS#7填充的数据
func unpad(buf []byte) []byte {
	length := len(buf)
	if length == 0 {
		return buf
	}
	unpadding := int(buf[length-1])
	return buf[:length-unpadding]
}