package main

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"testing"
	"time"

	"github.com/yuanyuanxiang/SimpleRemoter/server/go/protocol"
)

// connAuthTestKey is an arbitrary non-secret key used only to exercise the
// verifyConnAuth code path in tests. The production key is never compiled in:
// it is read from the YAMA_CONN_AUTH_KEY environment variable at startup.
const connAuthTestKey = "test-conn-auth-key"

// signConnAuth reproduces the client's ConnAuth signature: hex(HMAC-SHA256(key,
// clientID||timestamp||nonce)) as 64 lowercase hex chars. Mirrors
// client/IOCPClient.cpp PerformConnAuth + license.cpp signMessage.
func signConnAuth(key string, clientID, ts uint64, nonce [16]byte) string {
	mac := hmac.New(sha256.New, []byte(key))
	var msg [32]byte
	binary.LittleEndian.PutUint64(msg[0:8], clientID)
	binary.LittleEndian.PutUint64(msg[8:16], ts)
	copy(msg[16:32], nonce[:])
	mac.Write(msg[:])
	return hex.EncodeToString(mac.Sum(nil))
}

func testNonce() [16]byte {
	var n [16]byte
	for i := range n {
		n[i] = byte(i)
	}
	return n
}

// TestConnAuthHmacAlgorithmVector pins the signing algorithm and byte layout
// (HMAC-SHA256 over clientID||timestamp||nonce, lowercase hex) against a vector
// computed independently with Python's hmac/hashlib (not Go's crypto). The key
// value is irrelevant here — what matters is that the server reproduces the
// exact signing the client performs. The real key is supplied at runtime via
// YAMA_CONN_AUTH_KEY, never compiled in.
func TestConnAuthHmacAlgorithmVector(t *testing.T) {
	const want = "e59bf1259c77e065afa604cb5b1794b0c8c03fbb7e3473f5c3c34c2b94a90406"
	if got := signConnAuth(connAuthTestKey, 0x1122334455667788, 1700000000, testNonce()); got != want {
		t.Fatalf("HMAC mismatch:\n got  %s\n want %s", got, want)
	}
}

func TestVerifyConnAuth(t *testing.T) {
	h := &MyHandler{connAuthKey: []byte(connAuthTestKey)}

	build := func(ts uint64, sig []byte) []byte {
		pkt := make([]byte, protocol.ConnAuthPacketSize)
		pkt[0] = protocol.TokenConnAuth
		binary.LittleEndian.PutUint64(pkt[protocol.ConnAuthOffClientID:], 0x1122334455667788)
		binary.LittleEndian.PutUint64(pkt[protocol.ConnAuthOffTimestamp:], ts)
		for i := range 16 {
			pkt[protocol.ConnAuthOffNonce+i] = byte(i)
		}
		copy(pkt[protocol.ConnAuthOffSignature:], sig)
		return pkt
	}

	nonce := testNonce()

	t.Run("ok", func(t *testing.T) {
		ts := uint64(time.Now().Unix())
		sig := signConnAuth(connAuthTestKey, 0x1122334455667788, ts, nonce)
		if got := h.verifyConnAuth(build(ts, []byte(sig))); got != protocol.ConnAuthStatusOK {
			t.Fatalf("expected OK, got status %d", got)
		}
	})

	t.Run("bad_signature", func(t *testing.T) {
		ts := uint64(time.Now().Unix())
		sig := signConnAuth(connAuthTestKey, 0x1122334455667788, ts, nonce)
		pkt := build(ts, []byte(sig))
		pkt[protocol.ConnAuthOffSignature] ^= 0xff
		if got := h.verifyConnAuth(pkt); got != protocol.ConnAuthStatusBadSignature {
			t.Fatalf("expected BadSignature, got %d", got)
		}
	})

	t.Run("clock_skew", func(t *testing.T) {
		// Just outside the 300 s tolerance, with a signature valid for that
		// timestamp, so the skew check fires before the signature check.
		ts := uint64(time.Now().Unix()) - protocol.ConnAuthTimestampToleranceSec - 1
		sig := signConnAuth(connAuthTestKey, 0x1122334455667788, ts, nonce)
		if got := h.verifyConnAuth(build(ts, []byte(sig))); got != protocol.ConnAuthStatusClockSkew {
			t.Fatalf("expected ClockSkew, got %d", got)
		}
	})

	t.Run("bad_size", func(t *testing.T) {
		if got := h.verifyConnAuth([]byte{protocol.TokenConnAuth}); got != protocol.ConnAuthStatusBadSize {
			t.Fatalf("expected BadSize, got %d", got)
		}
	})
}