package main

import (
	"bytes"
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"os"
)

func EncryptAES(plainText, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	// PKCS#7 padding
	padLen := aes.BlockSize - len(plainText)%aes.BlockSize
	padding := bytes.Repeat([]byte{byte(padLen)}, padLen)
	plainText = append(plainText, padding...)

	cipherText := make([]byte, aes.BlockSize+len(plainText))
	iv := cipherText[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return nil, err
	}

	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(cipherText[aes.BlockSize:], plainText)

	return cipherText, nil
}

// DecryptAES decrypts the ciphertext using AES algorithm with the given key.
// Returns plaintext or an error if decryption fails.
func DecryptAES(cipherText, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	if len(cipherText) < aes.BlockSize {
		return nil, errors.New("ciphertext too short")
	}

	iv := cipherText[:aes.BlockSize]
	cipherText = cipherText[aes.BlockSize:]

	if len(cipherText)%aes.BlockSize != 0 {
		return nil, errors.New("ciphertext is not a multiple of the block size")
	}

	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(cipherText, cipherText)

	// PKCS#7 unpadding
	padLen := int(cipherText[len(cipherText)-1])
	if padLen < 1 || padLen > aes.BlockSize {
		return nil, errors.New("invalid padding")
	}
	cipherText = cipherText[:len(cipherText)-padLen]

	return cipherText, nil
}

// ReadPublicKeyFromFile 从文件中读取 PEM 格式的公钥
func ReadPublicKeyFromFile(publicKeyPath string) (*rsa.PublicKey, error) {
	publicKeyBytes, err := os.ReadFile(publicKeyPath)
	if err != nil {
		return nil, err
	}

	publicKeyPEM, _ := pem.Decode(publicKeyBytes)
	if publicKeyPEM == nil {
		return nil, errors.New("failed to decode PEM block containing public key")
	}

	publicKey, err := x509.ParsePKCS1PublicKey(publicKeyPEM.Bytes)
	if err != nil {
		return nil, err
	}

	return publicKey, nil
}

// ReadPrivateKeyFromFile 从文件中读取 PEM 格式的私钥
func ReadPrivateKeyFromFile(privateKeyPath string) (*rsa.PrivateKey, error) {
	privateKeyBytes, err := os.ReadFile(privateKeyPath)
	if err != nil {
		return nil, err
	}

	privateKeyPEM, _ := pem.Decode(privateKeyBytes)
	if privateKeyPEM == nil {
		return nil, errors.New("failed to decode PEM block containing private key")
	}

	privateKey, err := x509.ParsePKCS1PrivateKey(privateKeyPEM.Bytes)
	if err != nil {
		return nil, err
	}

	return privateKey, nil
}

type SignedEnvelope struct {
	Data      Envelope `json:"data"`
	Signature string   `json:"sign"`
}

// Sign 私钥签名
func (s *SignedEnvelope) Sign(privateKeyPath string) error {
	privateKey, err := ReadPrivateKeyFromFile(privateKeyPath)
	if err != nil {
		return err
	}
	marshal, err := json.Marshal(s.Data)
	if err != nil {
		return err
	}
	hashed := sha256.Sum256(marshal)
	signature, err := rsa.SignPKCS1v15(rand.Reader, privateKey, crypto.SHA256, hashed[:])
	if err != nil {
		panic(err)
	}
	s.Signature = base64.StdEncoding.EncodeToString(signature)
	return nil
}

// Verify 公钥验证
func (s *SignedEnvelope) Verify(publicKeyPath string) error {

	publicKey, err := ReadPublicKeyFromFile(publicKeyPath)
	if err != nil {
		return err
	}

	marshal, err := json.Marshal(s.Data)
	if err != nil {
		return err
	}
	hashed := sha256.Sum256(marshal)
	signature, err := base64.StdEncoding.DecodeString(s.Signature)
	if err != nil {
		return err
	}
	err = rsa.VerifyPKCS1v15(publicKey, crypto.SHA256, hashed[:], signature)
	if err != nil {
		return err
	}
	return nil
}

type Envelope struct {
	Data   string `json:"data"`
	Secret string `json:"secret"`
}

func (m *Envelope) SetEncryptData(value interface{}, publicKeyPath string) error {
	type Data struct {
		Data interface{} `json:"data"`
	}
	d := Data{
		value,
	}
	marshal, err := json.Marshal(&d)
	if err != nil {
		return err
	}
	// 生成随机的对称密钥
	symmetricKey := make([]byte, 32)
	if _, err := rand.Read(symmetricKey); err != nil {
		panic(err)
	}
	encryptData, err := EncryptAES(marshal, symmetricKey)
	if err != nil {
		return err
	}
	m.Data = base64.StdEncoding.EncodeToString(encryptData)

	publicKey, err := ReadPublicKeyFromFile(publicKeyPath)
	if err != nil {
		return err
	}
	encryptedSymmetricKey, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, symmetricKey)
	if err != nil {
		panic(err)
	}
	m.Secret = base64.StdEncoding.EncodeToString(encryptedSymmetricKey)
	return nil
}
func (m *Envelope) GetDecryptData(valueOut interface{}, privateKeyPath string) error {
	encryptData, err := base64.StdEncoding.DecodeString(m.Data)
	if err != nil {
		return err
	}
	encryptSecret, err := base64.StdEncoding.DecodeString(m.Secret)
	if err != nil {
		return err
	}

	privateKey, err := ReadPrivateKeyFromFile(privateKeyPath)
	if err != nil {
		return err
	}
	secret, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, encryptSecret)
	if err != nil {
		panic(err)
	}

	decryptData, err := DecryptAES(encryptData, secret)
	if err != nil {
		return err
	}
	type Data struct {
		Data interface{} `json:"data"`
	}
	d := Data{
		valueOut,
	}
	err = json.Unmarshal(decryptData, &d)
	if err != nil {
		return nil
	}
	return nil
}

// RSASend
//
//	@Description: 自动生成AES对称密钥，加密data得到密文，再使用接收者公钥加密对称密钥，最后使用发送者私钥对加密后的对称密钥和密文签名
//	@param data		需要加密的数据
//	@param senderPrivateKeyPath 发送者私钥
//	@param ReceiverPublicKeyPath 接收者公钥
//	@return SignedEnvelope 加密并签名的数据，序列化后发送即可
//	@return error 		错误
func RSASend(data interface{}, senderPrivateKeyPath, ReceiverPublicKeyPath string) (SignedEnvelope, error) {
	se := SignedEnvelope{}
	err := se.Data.SetEncryptData(data, ReceiverPublicKeyPath)
	if err != nil {
		return se, err
	}
	err = se.Sign(senderPrivateKeyPath)
	if err != nil {
		return se, err
	}
	return se, nil
}

// RSAReceive
//
//	@Description: 使用发送者公钥验证签名，再用接收者私钥解密对称密钥，最后用对称密钥解密密文
//	@param valueOut 传入指针，解密后的数据通过这个返回
//	@param signedEnvelope 加密并签名的数据，反序列化后得到
//	@param senderPublicKeyPath 发送者公钥
//	@param ReceiverPrivateKeyPath 接收者私钥
//	@return error 		错误
func RSAReceive(valueOut interface{}, signedEnvelope SignedEnvelope, senderPublicKeyPath, ReceiverPrivateKeyPath string) error {
	err := signedEnvelope.Verify(senderPublicKeyPath)
	if err != nil {
		return fmt.Errorf("签名校验失败")
	}
	err = signedEnvelope.Data.GetDecryptData(valueOut, ReceiverPrivateKeyPath)
	if err != nil {
		return fmt.Errorf("解密失败")
	}
	return nil
}

func Generate(publicKeyPath, privateKeyPath string) {
	// 生成 2048 位 RSA 密钥对
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		panic(err)
	}

	// 将私钥保存为 PKCS#1 格式 PEM 文件
	privateKeyBlock := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
	}
	privateKeyPem := pem.EncodeToMemory(privateKeyBlock)
	err = os.WriteFile(privateKeyPath, privateKeyPem, 0644)
	if err != nil {
		panic(err)
	}

	// 从私钥中提取公钥
	publicKey := &privateKey.PublicKey

	// 将公钥保存为 PKCS#1 格式 PEM 文件
	publicKeyBytes := x509.MarshalPKCS1PublicKey(publicKey)
	publicKeyBlock := &pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: publicKeyBytes,
	}
	publicKeyPem := pem.EncodeToMemory(publicKeyBlock)
	err = os.WriteFile(publicKeyPath, publicKeyPem, 0644)
	if err != nil {
		panic(err)
	}
}