package aerotun

import (
	"bytes"
	"encoding/binary"
	"errors"
	"testing"
)

const testPSK = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"

func testProtocol(t *testing.T) *protocolState {
	t.Helper()
	p, err := newProtocol(testPSK)
	if err != nil {
		t.Fatal(err)
	}
	return p
}
func testSeal(t *testing.T, p *protocolState, flags byte, b []byte, target sessionID) []byte {
	t.Helper()
	if flags != 0 && flags != FlagCompressedLZ4 && b == nil {
		b = []byte{5, 140}
	} // 1420-byte MTU advertisement
	out, err := p.seal(flags, b, target)
	if err != nil {
		t.Fatal(err)
	}
	return out
}
func establish(t *testing.T, a, b *protocolState) {
	t.Helper()
	hello := testSeal(t, a, flagHello, nil, zeroSession)
	_, _, sender, _, err := b.open(hello)
	if err != nil {
		t.Fatal(err)
	}
	ack := testSeal(t, b, flagAck, nil, sender)
	_, _, sender, _, err = a.open(ack)
	if err != nil {
		t.Fatal(err)
	}
	confirm := testSeal(t, a, flagConfirm, nil, sender)
	if _, _, _, _, err = b.open(confirm); err != nil {
		t.Fatal(err)
	}
}
func TestProtocolAuthenticationReplayAndRestart(t *testing.T) {
	a, b := testProtocol(t), testProtocol(t)
	establish(t, a, b)
	forward := testSeal(t, a, 0, []byte("forward"), zeroSession)
	reverse := testSeal(t, b, 0, []byte("reverse"), zeroSession)
	if bytes.Equal(forward[4:44], reverse[4:44]) {
		t.Fatal("directions reuse nonce/key context")
	}
	bad := append([]byte(nil), forward...)
	binary.BigEndian.PutUint64(bad[36:44], 1<<60)
	if _, _, _, _, err := b.open(bad); err == nil {
		t.Fatal("tampered counter authenticated")
	}
	got, _, _, _, err := b.open(forward)
	if err != nil || string(got) != "forward" {
		t.Fatalf("valid data rejected after forgery: %q %v", got, err)
	}
	if _, _, _, _, err = b.open(forward); !errors.Is(err, errReplay) {
		t.Fatalf("duplicate accepted: %v", err)
	}
	if _, _, _, _, err = a.open(forward); err == nil {
		t.Fatal("reflected data accepted")
	}
	if got, _, _, _, err = a.open(reverse); err != nil || string(got) != "reverse" {
		t.Fatal(err)
	}
	oldAck := testSeal(t, a, flagAck, nil, b.local)
	// Reboot sender: old packets and old ACKs cannot reactivate a retired session.
	a2 := testProtocol(t)
	establish(t, a2, b)
	if _, _, _, _, err = b.open(oldAck); err == nil {
		t.Fatal("retired peer session accepted")
	}
	if _, _, _, _, err = b.open(forward); err == nil {
		t.Fatal("old data accepted after sender restart")
	}
	fresh := testSeal(t, a2, 0, []byte("after restart"), zeroSession)
	if _, _, _, _, err = b.open(fresh); err != nil {
		t.Fatal(err)
	}
	// Reboot receiver: even authentic recordings target its obsolete boot ID.
	b2 := testProtocol(t)
	if _, _, _, _, err = b2.open(fresh); err == nil {
		t.Fatal("old data accepted after receiver restart")
	}
	establish(t, a2, b2)
	if _, _, _, _, err = b2.open(testSeal(t, a2, 0, []byte("new receiver"), zeroSession)); err != nil {
		t.Fatal(err)
	}
}
func TestHelloDoesNotReplaceSession(t *testing.T) {
	a, b := testProtocol(t), testProtocol(t)
	establish(t, a, b)
	stranger := testProtocol(t)
	_, _, _, _, err := b.open(testSeal(t, stranger, flagHello, nil, zeroSession))
	if err != nil {
		t.Fatal(err)
	}
	if b.remote != a.local {
		t.Fatal("HELLO replaced established session")
	}
	wrong, err := newProtocol("a completely different shared secret")
	if err != nil {
		t.Fatal(err)
	}
	if _, _, _, _, err = b.open(testSeal(t, wrong, flagHello, nil, zeroSession)); err == nil {
		t.Fatal("wrong PSK accepted")
	}
}
func TestReplayWindow(t *testing.T) {
	var w replayWindow
	for _, n := range []uint64{1, 3, 2, 66, 65} {
		if !w.accept(n) {
			t.Fatalf("valid counter %d rejected", n)
		}
	}
	for _, n := range []uint64{0, 1, 2, 3, 65, 66} {
		if w.accept(n) {
			t.Fatalf("invalid counter %d accepted", n)
		}
	}
}
func TestHandshakeConfirmationRetry(t *testing.T) {
	a, b := testProtocol(t), testProtocol(t)
	hello := testSeal(t, a, flagHello, nil, zeroSession)
	_, _, sender, _, err := b.open(hello)
	if err != nil {
		t.Fatal(err)
	}
	ack := testSeal(t, b, flagAck, nil, sender)
	_, _, _, _, err = a.open(ack)
	if err != nil {
		t.Fatal(err)
	}
	// First confirmation is lost. Repeated HELLO -> ACK -> CONFIRM still works.
	establish(t, a, b)
	if _, _, _, _, err = b.open(testSeal(t, a, 0, []byte("data"), zeroSession)); err != nil {
		t.Fatal(err)
	}
}