package bangcle

// Expected values come from pybyd 0.0.75 — regenerate with testdata/gen_fixtures.py.

import (
	"strings"
	"testing"
)

func TestRoundTrip(t *testing.T) {
	inputs := []string{
		"",
		"a",
		"0123456789abcdef", // exactly one block: must gain a whole pad block
		strings.Repeat("x", 4096),
		`{"request":"hello world"}`,
		"ünïcodé ✓",
	}
	for _, in := range inputs {
		out, err := Decode(Encode([]byte(in)))
		if err != nil {
			t.Fatalf("Decode(Encode(%.20q)): %v", in, err)
		}
		if string(out) != in {
			t.Errorf("round trip %.20q -> %.20q", in, out)
		}
	}
}

// TestVectors pins the white-box AES port against ciphertext from pybyd. If the
// block functions drift by a single index this fails; round-tripping alone
// would not, since it would drift symmetrically.
func TestVectors(t *testing.T) {
	vectors := []struct{ plain, want string }{
		{"", "FlpybH7xiohrkjagBOJY5Mg=="},
		{"hello", "FsFbBBn9/qw4cYEF9S1Jk6w=="},
		{"ünïcodé ✓", "FRqewj1Ypq1V2n8VqsiZPUA=="},
		{`{"functionType":"pwdLogin"}`, "FCHhJ73+HviiIT3dTDBeITnzHp5caNe/C50GlE8P0dqQ="},
	}
	for _, v := range vectors {
		if got := Encode([]byte(v.plain)); got != v.want {
			t.Errorf("Encode(%q)\n got %s\nwant %s", v.plain, got, v.want)
		}
	}
}

func TestDecodeRejectsMalformed(t *testing.T) {
	cases := map[string]string{
		"empty":          "",
		"missing prefix": "aGVsbG8=",
		"bad base64":     "F!!!!",
		"short block":    "F" + "AAAA", // 3 bytes: not a multiple of 16
	}
	for name, in := range cases {
		if _, err := Decode(in); err == nil {
			t.Errorf("%s: expected an error for %q", name, in)
		}
	}
}

// TestDecodeNormalises covers the url-safe alphabet, stray whitespace and the
// missing padding that BYD responses sometimes arrive with.
func TestDecodeNormalises(t *testing.T) {
	want := "the quick brown fox"
	env := Encode([]byte(want))

	variants := []string{
		env,
		strings.TrimRight(env, "="),
		strings.NewReplacer("+", "-", "/", "_").Replace(env),
		" " + env + "\n",
	}
	for _, v := range variants {
		got, err := Decode(v)
		if err != nil {
			t.Errorf("Decode(%q): %v", v, err)
			continue
		}
		if string(got) != want {
			t.Errorf("Decode(%q) = %q, want %q", v, got, want)
		}
	}
}

func TestStripPKCS7LeavesUnpaddedData(t *testing.T) {
	// BYD sometimes returns payloads without padding; pybyd passes them
	// through rather than erroring, and so must this.
	for _, in := range []string{"", "no padding here", "\x05\x05\x05"} {
		if got := StripPKCS7([]byte(in)); string(got) != in {
			t.Errorf("StripPKCS7(%q) = %q, want it unchanged", in, got)
		}
	}
	padded := append([]byte("abc"), 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13)
	if got := StripPKCS7(padded); string(got) != "abc" {
		t.Errorf("StripPKCS7(valid) = %q, want %q", got, "abc")
	}
}

func TestPadPKCS7AlwaysAddsAtLeastOneByte(t *testing.T) {
	for _, n := range []int{0, 1, 15, 16, 17, 32} {
		got := PadPKCS7(make([]byte, n))
		if len(got)%16 != 0 || len(got) <= n {
			t.Errorf("PadPKCS7(%d bytes) produced %d bytes", n, len(got))
		}
	}
}

func TestTablesLoaded(t *testing.T) {
	for _, spec := range tableSpecs {
		var got int
		switch spec.name {
		case "inv_round":
			got = len(tables.invRound)
		case "inv_xor":
			got = len(tables.invXor)
		case "inv_first":
			got = len(tables.invFirst)
		case "round":
			got = len(tables.round)
		case "xor":
			got = len(tables.xor)
		case "final":
			got = len(tables.final)
		case "perm_decrypt":
			got = len(tables.permDecrypt)
		case "perm_encrypt":
			got = len(tables.permEncrypt)
		}
		if got != spec.size {
			t.Errorf("table %s: %d bytes, want %d", spec.name, got, spec.size)
		}
	}
}

func TestLoadTablesRejectsBadContainer(t *testing.T) {
	cases := map[string][]byte{
		"too short":   []byte("BGTB"),
		"bad magic":   append([]byte("XXXX"), make([]byte, 100)...),
		"bad version": append([]byte("BGTB\x02\x00\x08\x00"), make([]byte, 100)...),
		"bad count":   append([]byte("BGTB\x01\x00\x03\x00"), make([]byte, 100)...),
	}
	for name, data := range cases {
		if _, err := loadTables(data); err == nil {
			t.Errorf("%s: expected an error", name)
		}
	}
}

func BenchmarkEncode(b *testing.B) {
	payload := []byte(strings.Repeat(`{"field":"value"},`, 100))
	b.SetBytes(int64(len(payload)))
	for i := 0; i < b.N; i++ {
		Encode(payload)
	}
}