package crypto

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"strings"
	"testing"
)

const (
	tok1 = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa="
	tok2 = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb="
	tok3 = "ccccccccccccccccccccccccccccccccccccccccccc="
)

func mustSecret(t *testing.T) *TunnelSecret {
	t.Helper()
	s, err := DeriveTunnelSecret(tok1, tok2, tok3)
	if err != nil {
		t.Fatalf("DeriveTunnelSecret: %v", err)
	}
	return s
}

func TestDeriveTunnelSecretIsDeterministic(t *testing.T) {
	a := mustSecret(t)
	b := mustSecret(t)

	if !bytes.Equal(a.master, b.master) {
		t.Error("same tokens produced different secrets")
	}

	other, err := DeriveTunnelSecret(tok1, tok2, "different")
	if err != nil {
		t.Fatalf("DeriveTunnelSecret: %v", err)
	}
	if bytes.Equal(a.master, other.master) {
		t.Error("different tokens produced the same secret")
	}
}

func TestDeriveTunnelSecretRequiresAllTokens(t *testing.T) {
	for _, tc := range [][3]string{
		{"", tok2, tok3},
		{tok1, "", tok3},
		{tok1, tok2, ""},
	} {
		if _, err := DeriveTunnelSecret(tc[0], tc[1], tc[2]); err == nil {
			t.Errorf("DeriveTunnelSecret(%q, %q, %q) succeeded, want an error", tc[0], tc[1], tc[2])
		}
	}
}

// The whole point of the design: the relay knows the group hash, so the
// tunnel secret must not be derivable from it.
func TestTunnelSecretDiffersFromGroupHash(t *testing.T) {
	s := mustSecret(t)

	combined := strings.Join([]string{tok1, tok2, tok3}, ":")
	sum := sha256.Sum256([]byte(combined))
	groupHash := hex.EncodeToString(sum[:])

	if hex.EncodeToString(s.master) == groupHash {
		t.Fatal("tunnel secret equals the group hash — the relay could decrypt everything")
	}

	key, err := s.TunnelKey("tunnel-1")
	if err != nil {
		t.Fatalf("TunnelKey: %v", err)
	}
	if hex.EncodeToString(key) == groupHash {
		t.Fatal("tunnel key equals the group hash")
	}
	if bytes.Equal(key, s.master) {
		t.Error("tunnel key equals the master secret; it should be derived per tunnel")
	}
}

func TestTunnelKeyIsPerTunnel(t *testing.T) {
	s := mustSecret(t)

	a, _ := s.TunnelKey("tunnel-a")
	b, _ := s.TunnelKey("tunnel-b")
	aAgain, _ := s.TunnelKey("tunnel-a")

	if bytes.Equal(a, b) {
		t.Error("different tunnel IDs produced the same key")
	}
	if !bytes.Equal(a, aAgain) {
		t.Error("same tunnel ID produced different keys")
	}
	if len(a) != keySize {
		t.Errorf("key is %d bytes, want %d", len(a), keySize)
	}
}

func TestSealOpenRoundTrip(t *testing.T) {
	s := mustSecret(t)
	key, _ := s.TunnelKey("t1")

	sealer, err := NewSealer(key, DirShareToUse)
	if err != nil {
		t.Fatalf("NewSealer: %v", err)
	}
	opener, err := NewOpener(key, DirShareToUse)
	if err != nil {
		t.Fatalf("NewOpener: %v", err)
	}

	messages := [][]byte{
		[]byte("first"),
		[]byte(""),
		bytes.Repeat([]byte{0xAB}, 65536),
		[]byte("last"),
	}

	for i, want := range messages {
		frame, err := sealer.Seal(want)
		if err != nil {
			t.Fatalf("Seal %d: %v", i, err)
		}
		if len(frame) != len(want)+FrameOverhead {
			t.Errorf("frame %d is %d bytes, want %d", i, len(frame), len(want)+FrameOverhead)
		}
		// The plaintext must not be visible on the wire.
		if len(want) > 8 && bytes.Contains(frame, want) {
			t.Errorf("frame %d contains its plaintext", i)
		}

		got, err := opener.Open(frame)
		if err != nil {
			t.Fatalf("Open %d: %v", i, err)
		}
		if !bytes.Equal(got, want) {
			t.Errorf("frame %d round-tripped to %q, want %q", i, got, want)
		}
	}
}

// Both ends derive the same tunnel key, so the direction byte is the only
// thing keeping their nonce spaces apart.
func TestDirectionsUseSeparateNonceSpaces(t *testing.T) {
	s := mustSecret(t)
	key, _ := s.TunnelKey("t1")

	shareToUse, _ := NewSealer(key, DirShareToUse)
	useToShare, _ := NewSealer(key, DirUseToShare)

	plaintext := []byte("identical plaintext")
	a, _ := shareToUse.Seal(plaintext)
	b, _ := useToShare.Seal(plaintext)

	if bytes.Equal(a, b) {
		t.Fatal("both directions produced identical ciphertext — nonce reuse")
	}
	// Same counter, so any difference must come from the direction byte.
	if !bytes.Equal(a[:counterSize], b[:counterSize]) {
		t.Fatal("test assumption broken: counters differ")
	}

	// A frame from one direction must not open with the other direction's opener.
	wrongWay, _ := NewOpener(key, DirUseToShare)
	if _, err := wrongWay.Open(a); err == nil {
		t.Error("a frame opened under the wrong direction")
	}
}

func TestOpenRejectsTampering(t *testing.T) {
	s := mustSecret(t)
	key, _ := s.TunnelKey("t1")
	sealer, _ := NewSealer(key, DirShareToUse)

	original, _ := sealer.Seal([]byte("sensitive usb traffic"))

	tests := []struct {
		name   string
		mutate func([]byte) []byte
	}{
		{"flipped ciphertext bit", func(f []byte) []byte {
			f[counterSize+2] ^= 0x01
			return f
		}},
		{"flipped counter bit", func(f []byte) []byte {
			f[0] ^= 0x80
			return f
		}},
		{"truncated tag", func(f []byte) []byte { return f[:len(f)-1] }},
		{"appended byte", func(f []byte) []byte { return append(f, 0x00) }},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			frame := append([]byte(nil), original...)
			opener, _ := NewOpener(key, DirShareToUse)
			if _, err := opener.Open(tt.mutate(frame)); err == nil {
				t.Error("tampered frame was accepted")
			}
		})
	}
}

func TestOpenRejectsWrongKey(t *testing.T) {
	s := mustSecret(t)
	good, _ := s.TunnelKey("t1")
	bad, _ := s.TunnelKey("t2")

	sealer, _ := NewSealer(good, DirShareToUse)
	frame, _ := sealer.Seal([]byte("secret"))

	opener, _ := NewOpener(bad, DirShareToUse)
	if _, err := opener.Open(frame); err == nil {
		t.Error("frame opened under a key from a different tunnel")
	}
}

func TestOpenRejectsReplay(t *testing.T) {
	s := mustSecret(t)
	key, _ := s.TunnelKey("t1")

	sealer, _ := NewSealer(key, DirShareToUse)
	opener, _ := NewOpener(key, DirShareToUse)

	frame, _ := sealer.Seal([]byte("do this once"))

	if _, err := opener.Open(append([]byte(nil), frame...)); err != nil {
		t.Fatalf("first delivery: %v", err)
	}
	if _, err := opener.Open(append([]byte(nil), frame...)); !errors.Is(err, ErrReplay) {
		t.Errorf("replayed frame gave %v, want ErrReplay", err)
	}
}

// A forged frame carrying a huge counter must not poison the replay window
// and lock out the genuine frames that follow.
func TestForgedFrameDoesNotAdvanceCounter(t *testing.T) {
	s := mustSecret(t)
	key, _ := s.TunnelKey("t1")

	sealer, _ := NewSealer(key, DirShareToUse)
	opener, _ := NewOpener(key, DirShareToUse)

	forged := make([]byte, FrameOverhead+4)
	for i := range forged[:counterSize] {
		forged[i] = 0xFF
	}
	if _, err := opener.Open(forged); err == nil {
		t.Fatal("forged frame was accepted")
	}

	genuine, _ := sealer.Seal([]byte("real traffic"))
	got, err := opener.Open(genuine)
	if err != nil {
		t.Fatalf("genuine frame rejected after a forgery: %v", err)
	}
	if string(got) != "real traffic" {
		t.Errorf("got %q", got)
	}
}

func TestOpenRejectsUndersizedFrame(t *testing.T) {
	s := mustSecret(t)
	key, _ := s.TunnelKey("t1")
	opener, _ := NewOpener(key, DirShareToUse)

	for _, size := range []int{0, 1, counterSize, FrameOverhead - 1} {
		if _, err := opener.Open(make([]byte, size)); err == nil {
			t.Errorf("frame of %d bytes was accepted", size)
		}
	}
}

func TestPeerTokenBindsToContext(t *testing.T) {
	s := mustSecret(t)

	a := s.PeerToken("tunnel-1")
	b := s.PeerToken("tunnel-2")
	aAgain := s.PeerToken("tunnel-1")

	if a == b {
		t.Error("different contexts produced the same token")
	}
	if a != aAgain {
		t.Error("same context produced different tokens")
	}
	if len(a) != 64 {
		t.Errorf("token is %d hex chars, want 64", len(a))
	}

	// A different group must not be able to produce a matching token.
	other, _ := DeriveTunnelSecret(tok1, tok2, "different")
	if other.PeerToken("tunnel-1") == a {
		t.Error("a different group secret produced the same peer token")
	}
}