aice/utils/wechat_pay.go

package utils

import (
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"
)

type WxPayClient struct {
	MchID      string
	SerialNo   string
	NotifyURL  string
	JSAPIAppID string
	H5AppID    string
	AppAppID   string
	PrivateKey *rsa.PrivateKey
	HTTPClient *http.Client
	APIBase    string
	Enabled    bool
	UserAgent  string
}

type WxPayAmount struct {
	Total    int    `json:"total"`
	Currency string `json:"currency,omitempty"`
}

type WxPayPayer struct {
	OpenID string `json:"openid,omitempty"`
}

type WxPayH5Info struct {
	Type string `json:"type"`
}

type WxPaySceneInfo struct {
	PayerClientIP string       `json:"payer_client_ip,omitempty"`
	H5Info        *WxPayH5Info `json:"h5_info,omitempty"`
}

type wxPayTxnRequest struct {
	AppID       string          `json:"appid"`
	MchID       string          `json:"mchid"`
	Description string          `json:"description"`
	OutTradeNo  string          `json:"out_trade_no"`
	NotifyURL   string          `json:"notify_url"`
	Attach      string          `json:"attach,omitempty"`
	Amount      WxPayAmount     `json:"amount"`
	Payer       *WxPayPayer     `json:"payer,omitempty"`
	SceneInfo   *WxPaySceneInfo `json:"scene_info,omitempty"`
}

type wxPayTxnResp struct {
	PrepayID string `json:"prepay_id"`
	H5URL    string `json:"h5_url"`
}

type WxMiniProgramOrderReq struct {
	Description string
	OutTradeNo  string
	Total       int
	Currency    string
	OpenID      string
	NotifyURL   string
	Attach      string
	AppID       string
}

type WxMiniProgramPayResult struct {
	AppID     string `json:"appId"`
	TimeStamp string `json:"timeStamp"`
	NonceStr  string `json:"nonceStr"`
	Package   string `json:"package"`
	SignType  string `json:"signType"`
	PaySign   string `json:"paySign"`
	PrepayID  string `json:"prepayId"`
}

type WxH5OrderReq struct {
	Description   string
	OutTradeNo    string
	Total         int
	Currency      string
	NotifyURL     string
	Attach        string
	PayerClientIP string
	AppID         string
}

type WxH5PayResult struct {
	H5URL string `json:"h5Url"`
}

type WxAppOrderReq struct {
	Description string
	OutTradeNo  string
	Total       int
	Currency    string
	NotifyURL   string
	Attach      string
	AppID       string
}

type WxAppPayResult struct {
	AppID     string `json:"appid"`
	PartnerID string `json:"partnerid"`
	PrepayID  string `json:"prepayid"`
	Package   string `json:"package"`
	NonceStr  string `json:"noncestr"`
	TimeStamp string `json:"timestamp"`
	Sign      string `json:"sign"`
}

var wxPayClient *WxPayClient

func InitWxPay() error {
	enabled := redisBool("wxpay_enable", false)
	if !enabled {
		return nil
	}

	mchID := redisString("wxpay_mch_id", "")
	serialNo := redisString("wxpay_serial_no", "")
	privateKeyPath := redisString("wxpay_private_key_path", "")
	notifyURL := redisString("wxpay_notify_url", "")
	jsapiAppID := redisString("wxpay_jsapi_appid", "")
	h5AppID := redisString("wxpay_h5_appid", "")
	appAppID := redisString("wxpay_app_appid", "")
	apiBase := redisString("wxpay_api_base", "https://api.mch.weixin.qq.com")

	if mchID == "" || serialNo == "" || privateKeyPath == "" || notifyURL == "" {
		return errors.New("wxpay config missing: wxpay_mch_id/wxpay_serial_no/wxpay_private_key_path/wxpay_notify_url are required")
	}

	privateKey, err := loadRSAPrivateKey(privateKeyPath)
	if err != nil {
		return err
	}

	wxPayClient = &WxPayClient{
		MchID:      mchID,
		SerialNo:   serialNo,
		NotifyURL:  notifyURL,
		JSAPIAppID: jsapiAppID,
		H5AppID:    h5AppID,
		AppAppID:   appAppID,
		PrivateKey: privateKey,
		HTTPClient: &http.Client{Timeout: 15 * time.Second},
		APIBase:    strings.TrimRight(apiBase, "/"),
		Enabled:    true,
		UserAgent:  "think-go-wxpay/1.0",
	}

	return nil
}

func GetWxPayClient() (*WxPayClient, error) {
	if wxPayClient == nil || !wxPayClient.Enabled {
		return nil, errors.New("wxpay not initialized, set wxpay_enable=true and call InitWxPay")
	}
	return wxPayClient, nil
}

func CreateWxMiniProgramPay(req WxMiniProgramOrderReq) (*WxMiniProgramPayResult, error) {
	client, err := GetWxPayClient()
	if err != nil {
		return nil, err
	}
	return client.CreateMiniProgramPay(req)
}

func CreateWxH5Pay(req WxH5OrderReq) (*WxH5PayResult, error) {
	client, err := GetWxPayClient()
	if err != nil {
		return nil, err
	}
	return client.CreateH5Pay(req)
}

func CreateWxAppPay(req WxAppOrderReq) (*WxAppPayResult, error) {
	client, err := GetWxPayClient()
	if err != nil {
		return nil, err
	}
	return client.CreateAppPay(req)
}

func (c *WxPayClient) CreateMiniProgramPay(req WxMiniProgramOrderReq) (*WxMiniProgramPayResult, error) {
	appID := req.AppID
	if appID == "" {
		appID = c.JSAPIAppID
	}
	if appID == "" {
		return nil, errors.New("jsapi appid is required")
	}
	if req.OpenID == "" {
		return nil, errors.New("openid is required")
	}
	if req.OutTradeNo == "" || req.Description == "" || req.Total <= 0 {
		return nil, errors.New("description/outTradeNo/total are required")
	}

	requestBody := wxPayTxnRequest{
		AppID:       appID,
		MchID:       c.MchID,
		Description: req.Description,
		OutTradeNo:  req.OutTradeNo,
		NotifyURL:   firstNotEmpty(req.NotifyURL, c.NotifyURL),
		Attach:      req.Attach,
		Amount: WxPayAmount{
			Total:    req.Total,
			Currency: firstNotEmpty(req.Currency, "CNY"),
		},
		Payer: &WxPayPayer{OpenID: req.OpenID},
	}

	var resp wxPayTxnResp
	if err := c.postJSON("/v3/pay/transactions/jsapi", requestBody, &resp); err != nil {
		return nil, err
	}
	if resp.PrepayID == "" {
		return nil, errors.New("wxpay jsapi response missing prepay_id")
	}

	timeStamp := strconv.FormatInt(time.Now().Unix(), 10)
	nonceStr := randomString(32)
	pkg := "prepay_id=" + resp.PrepayID
	signPayload := appID + "\n" + timeStamp + "\n" + nonceStr + "\n" + pkg + "\n"
	paySign, err := c.signBase64(signPayload)
	if err != nil {
		return nil, err
	}

	return &WxMiniProgramPayResult{
		AppID:     appID,
		TimeStamp: timeStamp,
		NonceStr:  nonceStr,
		Package:   pkg,
		SignType:  "RSA",
		PaySign:   paySign,
		PrepayID:  resp.PrepayID,
	}, nil
}

func (c *WxPayClient) CreateH5Pay(req WxH5OrderReq) (*WxH5PayResult, error) {
	appID := req.AppID
	if appID == "" {
		appID = c.H5AppID
	}
	if appID == "" {
		return nil, errors.New("h5 appid is required")
	}
	if req.OutTradeNo == "" || req.Description == "" || req.Total <= 0 {
		return nil, errors.New("description/outTradeNo/total are required")
	}
	if req.PayerClientIP == "" {
		return nil, errors.New("payer_client_ip is required for h5 pay")
	}

	requestBody := wxPayTxnRequest{
		AppID:       appID,
		MchID:       c.MchID,
		Description: req.Description,
		OutTradeNo:  req.OutTradeNo,
		NotifyURL:   firstNotEmpty(req.NotifyURL, c.NotifyURL),
		Attach:      req.Attach,
		Amount: WxPayAmount{
			Total:    req.Total,
			Currency: firstNotEmpty(req.Currency, "CNY"),
		},
		SceneInfo: &WxPaySceneInfo{
			PayerClientIP: req.PayerClientIP,
			H5Info:        &WxPayH5Info{Type: "Wap"},
		},
	}

	var resp wxPayTxnResp
	if err := c.postJSON("/v3/pay/transactions/h5", requestBody, &resp); err != nil {
		return nil, err
	}
	if resp.H5URL == "" {
		return nil, errors.New("wxpay h5 response missing h5_url")
	}
	return &WxH5PayResult{H5URL: resp.H5URL}, nil
}

func (c *WxPayClient) CreateAppPay(req WxAppOrderReq) (*WxAppPayResult, error) {
	appID := req.AppID
	if appID == "" {
		appID = c.AppAppID
	}
	if appID == "" {
		return nil, errors.New("app appid is required")
	}
	if req.OutTradeNo == "" || req.Description == "" || req.Total <= 0 {
		return nil, errors.New("description/outTradeNo/total are required")
	}

	requestBody := wxPayTxnRequest{
		AppID:       appID,
		MchID:       c.MchID,
		Description: req.Description,
		OutTradeNo:  req.OutTradeNo,
		NotifyURL:   firstNotEmpty(req.NotifyURL, c.NotifyURL),
		Attach:      req.Attach,
		Amount: WxPayAmount{
			Total:    req.Total,
			Currency: firstNotEmpty(req.Currency, "CNY"),
		},
	}

	var resp wxPayTxnResp
	if err := c.postJSON("/v3/pay/transactions/app", requestBody, &resp); err != nil {
		return nil, err
	}
	if resp.PrepayID == "" {
		return nil, errors.New("wxpay app response missing prepay_id")
	}

	timeStamp := strconv.FormatInt(time.Now().Unix(), 10)
	nonceStr := randomString(32)
	signPayload := appID + "\n" + timeStamp + "\n" + nonceStr + "\n" + resp.PrepayID + "\n"
	sign, err := c.signBase64(signPayload)
	if err != nil {
		return nil, err
	}

	return &WxAppPayResult{
		AppID:     appID,
		PartnerID: c.MchID,
		PrepayID:  resp.PrepayID,
		Package:   "Sign=WXPay",
		NonceStr:  nonceStr,
		TimeStamp: timeStamp,
		Sign:      sign,
	}, nil
}

func (c *WxPayClient) postJSON(path string, payload interface{}, out interface{}) error {
	bodyBytes, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	fullURL := c.APIBase + path

	auth, err := c.buildAuthorization("POST", fullURL, string(bodyBytes))
	if err != nil {
		return err
	}

	req, err := http.NewRequest(http.MethodPost, fullURL, bytes.NewReader(bodyBytes))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Authorization", auth)
	req.Header.Set("User-Agent", c.UserAgent)

	resp, err := c.HTTPClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return fmt.Errorf("wxpay request failed: status=%d body=%s", resp.StatusCode, string(respBody))
	}
	if out == nil {
		return nil
	}
	if err := json.Unmarshal(respBody, out); err != nil {
		return fmt.Errorf("parse wxpay response failed: %w, body=%s", err, string(respBody))
	}
	return nil
}

func (c *WxPayClient) buildAuthorization(method, fullURL, body string) (string, error) {
	u, err := url.Parse(fullURL)
	if err != nil {
		return "", err
	}
	canonicalURL := u.Path
	if u.RawQuery != "" {
		canonicalURL += "?" + u.RawQuery
	}
	nonce := randomString(32)
	timestamp := strconv.FormatInt(time.Now().Unix(), 10)
	message := method + "\n" + canonicalURL + "\n" + timestamp + "\n" + nonce + "\n" + body + "\n"

	signature, err := c.signBase64(message)
	if err != nil {
		return "", err
	}
	token := fmt.Sprintf(`mchid="%s",nonce_str="%s",timestamp="%s",serial_no="%s",signature="%s"`,
		c.MchID, nonce, timestamp, c.SerialNo, signature,
	)
	return "WECHATPAY2-SHA256-RSA2048 " + token, nil
}

func (c *WxPayClient) signBase64(message string) (string, error) {
	sum := sha256.Sum256([]byte(message))
	sign, err := rsa.SignPKCS1v15(rand.Reader, c.PrivateKey, crypto.SHA256, sum[:])
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(sign), nil
}

func loadRSAPrivateKey(path string) (*rsa.PrivateKey, error) {
	raw, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	block, _ := pem.Decode(raw)
	if block == nil {
		return nil, errors.New("invalid private key pem")
	}

	if key, err := x509.ParsePKCS8PrivateKey(block.Bytes); err == nil {
		if rsaKey, ok := key.(*rsa.PrivateKey); ok {
			return rsaKey, nil
		}
	}
	if key, err := x509.ParsePKCS1PrivateKey(block.Bytes); err == nil {
		return key, nil
	}

	return nil, errors.New("unsupported private key format")
}

func randomString(n int) string {
	const chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	if n <= 0 {
		return ""
	}
	b := make([]byte, n)
	rb := make([]byte, n)
	if _, err := rand.Read(rb); err != nil {
		return strconv.FormatInt(time.Now().UnixNano(), 10)
	}
	for i := 0; i < n; i++ {
		b[i] = chars[int(rb[i])%len(chars)]
	}
	return string(b)
}

func firstNotEmpty(values ...string) string {
	for _, v := range values {
		if strings.TrimSpace(v) != "" {
			return strings.TrimSpace(v)
		}
	}
	return ""
}