package payments

import (
	"bytes"
	"context"
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"net/http"
	"strconv"
	"testing"
	"time"
)

func TestVerifiedNotificationRejectsReplayAndTampering(t *testing.T) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatal(err)
	}
	pubDER, err := x509.MarshalPKIXPublicKey(&key.PublicKey)
	if err != nil {
		t.Fatal(err)
	}
	privDER, err := x509.MarshalPKCS8PrivateKey(key)
	if err != nil {
		t.Fatal(err)
	}
	config := Config{MchID: "mch", AppID: "app", NotifyURL: "https://example.com/notify", MerchantSerial: "serial",
		PrivateKeyPEM:     string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: privDER})),
		WechatPublicKeyID: "PUB_KEY_ID_123", WechatPublicKeyPEM: string(pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: pubDER})),
		APIv3Key: "12345678901234567890123456789012"}
	provider, err := NewWeChatProvider(config)
	if err != nil {
		t.Fatal(err)
	}
	wp := provider.(*wechatProvider)
	txn := map[string]any{"out_trade_no": "order1", "mchid": "mch", "appid": "app", "transaction_id": "wx1",
		"trade_state": "SUCCESS", "amount": map[string]any{"total": 10000, "currency": "CNY"}, "success_time": time.Now().UTC().Format(time.RFC3339)}
	plain, _ := json.Marshal(txn)
	block, _ := aes.NewCipher([]byte(config.APIv3Key))
	aead, _ := cipher.NewGCM(block)
	nonce := "123456789012"
	ciphertext := aead.Seal(nil, []byte(nonce), plain, []byte("associated"))
	body, _ := json.Marshal(map[string]any{"id": "event1", "event_type": "TRANSACTION.SUCCESS", "resource_type": "encrypt-resource",
		"resource": map[string]any{"algorithm": "AEAD_AES_256_GCM", "ciphertext": base64.StdEncoding.EncodeToString(ciphertext), "associated_data": "associated", "nonce": nonce}})
	makeRequest := func(timestamp int64, tamper bool) *http.Request {
		ts := strconv.FormatInt(timestamp, 10)
		signed := fmt.Sprintf("%s\n%s\n%s\n", ts, "nonce1", body)
		hash := sha256.Sum256([]byte(signed))
		signature, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA256, hash[:])
		if err != nil {
			t.Fatal(err)
		}
		actualBody := body
		if tamper {
			actualBody = append(bytes.Clone(body), ' ')
		}
		req, _ := http.NewRequest(http.MethodPost, "https://example.com/notify", bytes.NewReader(actualBody))
		req.Header.Set("Wechatpay-Serial", config.WechatPublicKeyID)
		req.Header.Set("Wechatpay-Timestamp", ts)
		req.Header.Set("Wechatpay-Nonce", "nonce1")
		req.Header.Set("Wechatpay-Signature", base64.StdEncoding.EncodeToString(signature))
		return req
	}
	valid, err := wp.ParseNotification(context.Background(), makeRequest(time.Now().Unix(), false))
	if err != nil || valid.AmountFen != 10000 || valid.OrderID != "order1" {
		t.Fatalf("valid notification: %v %+v", err, valid)
	}
	if _, err := wp.ParseNotification(context.Background(), makeRequest(time.Now().Add(-10*time.Minute).Unix(), false)); err == nil {
		t.Fatal("accepted stale timestamp")
	}
	if _, err := wp.ParseNotification(context.Background(), makeRequest(time.Now().Unix(), true)); err == nil {
		t.Fatal("accepted tampered body")
	}
}