package rpc

import (
	"context"
	"crypto"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"eta/eta_bridge/global"
	"eta/eta_bridge/rpc/sso"
	"fmt"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"io/ioutil"
	"strconv"
)

type encryptedRequest struct {
	Message   []byte `json:"ciphertext"`
	Nonce     string `json:"nonce"`     // 添加随机字符串
	Timestamp int64  `json:"timestamp"` // 添加时间戳
}

// 自定义拦截器
func SignatureInterceptor(ctx context.Context, req interface{}, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
	// 获取签名
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		return nil, status.Errorf(codes.InvalidArgument, "metadata is missing")
	}

	// 从元数据中提取签名相关的信息
	nonceStr, ok := md["nonce"]
	if !ok || len(nonceStr) == 0 {
		return nil, status.Errorf(codes.InvalidArgument, "nonce不能为空")
	}

	timestampStr, ok := md["timestamp"]
	if !ok || len(timestampStr) == 0 {
		return nil, status.Errorf(codes.InvalidArgument, "时间戳不能为空")
	}
	timestamp, err := strconv.ParseInt(timestampStr[0], 10, 64)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "无效的时间戳: %v", err)
	}

	signature, ok := md["signature"]
	if !ok || len(signature) == 0 {
		return nil, status.Errorf(codes.InvalidArgument, "签名信息不能为空")
	}
	message := req.(*sso.LoginRequest)
	str, err := json.Marshal(message)
	decrytData := encryptedRequest{
		Message:   str,
		Nonce:     nonceStr[0], // 添加随机字符串
		Timestamp: timestamp,   // 添加时间戳
	}
	dds, err := json.Marshal(decrytData)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "json序列化失败")
	}
	publicKey, err := parsePublicKeyFromPEM([]byte(""))
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "公钥解析失败")
	}
	fmt.Println(string(dds))
	// 验证签名
	if !verifySignature(string(dds), signature[0], publicKey) {
		return nil, status.Errorf(codes.PermissionDenied, "invalid signature")
	}
	return handler(ctx, req)
}

// 验证签名
func verifySignature(message, signature string, publicKey *rsa.PublicKey) bool {
	hash := sha256.Sum256([]byte(message))
	signatureBytes, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		return false
	}
	err = rsa.VerifyPKCS1v15(publicKey, crypto.SHA256, hash[:], signatureBytes)
	return err == nil
}

func parsePublicKeyFromPEM(pemBytes []byte) (pubKey *rsa.PublicKey, err error) {
	pemBlock, err := ioutil.ReadFile("./config/rsa_public_key.pem")
	block, _ := pem.Decode(pemBlock)
	if block == nil {
		global.LOG.Error("公钥解析失败")
	}
	pubKey, err = x509.ParsePKCS1PublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	return
}