package utils

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"
	"github.com/tjfoc/gmsm/sm2"
	"github.com/tjfoc/gmsm/x509"
	"os"
)

type SM struct {
	PrivateFile string
	PublicFile  string
}

func (s *SM) CreateSm2Sig(msg []byte) ([]byte, *sm2.PublicKey, error) {
	//读取密钥对
	pri, err := s.readPemCxt(s.PrivateFile)
	privateKey, _ := x509.ReadPrivateKeyFromPem(pri, nil)
	c := sm2.P256Sm2() //椭圆曲线
	priKey := new(sm2.PrivateKey)
	priKey.PublicKey.Curve = c
	priKey.D = privateKey.D
	priKey.PublicKey.X = privateKey.X
	priKey.PublicKey.Y = privateKey.Y

	sign, err := priKey.Sign(rand.Reader, msg, nil) //sm2签名
	if err != nil {
		return nil, nil, err
	}
	return sign, &priKey.PublicKey, err
}

func (s *SM) readPemCxt(path string) ([]byte, error) {
	file, err := os.Open(path)
	if err != nil {
		return []byte{}, err
	}
	defer func(file *os.File) {
		err := file.Close()
		if err != nil {
			fmt.Println(err)
		}
	}(file)
	fileInfo, err := file.Stat()
	if err != nil {
		return []byte{}, err
	}
	buf := make([]byte, fileInfo.Size())
	_, err = file.Read(buf)
	if err != nil {
		return []byte{}, err
	}
	return buf, err
}

func (s *SM) Decrypt(data string) (string, error) {
	pri, err := s.readPemCxt(s.PrivateFile)
	if err != nil {
		return "", err
	}

	privateKeyFromPem, err := x509.ReadPrivateKeyFromPem(pri, nil)
	if err != nil {
		return "", err
	}

	cipherTxt, err := base64.StdEncoding.DecodeString(data)
	if err != nil {
		return "", err
	}
	plainTxt, err := privateKeyFromPem.DecryptAsn1(cipherTxt)
	if err != nil {
		return "", err
	}
	return string(plainTxt), nil
}

func (s *SM) Encrypt(data string) (string, error) {
	pub, err := s.readPemCxt(s.PublicFile)
	if err != nil {
		return "", err
	}
	//read public key
	publicKeyFromPem, err := x509.ReadPublicKeyFromPem(pub)
	if err != nil {
		fmt.Println(err)
		return "", err
	}
	cipherTxt, err := publicKeyFromPem.EncryptAsn1([]byte(data), rand.Reader)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(cipherTxt), nil
}

func (s *SM) VerSm2Sig(pub *sm2.PublicKey, msg []byte, sign []byte) bool {
	ok := pub.Verify(msg, sign)
	if !ok {
		return false
	}
	return true
}