package protocol_test

import (
	"testing"

	"git.n1ko.dev/Niko/niko_trust/pkg/protocol"
)

// Mutation tests: any single-byte change to the canonical bytes or to the
// signature must break verification, and every frozen vector must still
// verify under its own signature. This is the strongest form of the
// "the bytes that arrived are the bytes that are verified" property
// (docs/IMPLEMENTATION_NOTES.md property 2): if a verifier re-serialised the
// decoded object, a mutation that survived round-tripping would pass.

// verifyFor returns an error when the given vector no longer verifies.
func verifyFor(t *testing.T, name string, reqTce, reqSig, rotTce, rotSig, tceB, sig []byte) error {
	t.Helper()
	switch name {
	case "identity/nikocraft", "identity/niko":
		_, err := protocol.VerifyIdentity(tceB, sig)
		return err
	case "claim/boolean", "claim/all-value-types":
		_, err := protocol.VerifyClaim(tceB, sig)
		return err
	case "revocation/boolean-claim":
		_, err := protocol.VerifyRevocation(tceB, sig)
		return err
	case "approval_request/ban":
		_, err := protocol.VerifyApprovalRequest(tceB, sig)
		return err
	case "approval_response/allow", "approval_response/deny":
		_, err := protocol.VerifyApprovalResponse(reqTce, reqSig, tceB, sig)
		return err
	case "auth_assertion/ws":
		_, err := protocol.VerifyAuthAssertion(tceB, sig, "trust.n1ko.dev")
		return err
	case "delegation/minimal", "delegation/multi-predicate":
		_, err := protocol.VerifyDelegationClaim(tceB, sig)
		return err
	case "key_rotation/request":
		_, err := protocol.VerifyKeyRotationRequest(tceB, sig)
		return err
	case "key_rotation/confirm":
		_, err := protocol.VerifyKeyRotationConfirm(rotTce, rotSig, tceB, sig)
		return err
	}
	return nil
}

func TestMutationSweep(t *testing.T) {
	vf := loadVectors(t)
	byName := map[string]*vectorEntry{}
	for i := range vf.Vectors {
		byName[vf.Vectors[i].Name] = &vf.Vectors[i]
	}
	req := byName["approval_request/ban"]
	reqTce := mustHex(t, req.TCEHex)
	reqSig := mustHex(t, req.SignatureHex)
	rot := byName["key_rotation/request"]
	rotTce := mustHex(t, rot.TCEHex)
	rotSig := mustHex(t, rot.SignatureHex)

	for _, v := range vf.Vectors {
		t.Run(v.Name, func(t *testing.T) {
			b := mustHex(t, v.TCEHex)
			sig := mustHex(t, v.SignatureHex)

			// The unchanged vector must verify.
			if err := verifyFor(t, v.Name, reqTce, reqSig, rotTce, rotSig, b, sig); err != nil {
				t.Fatalf("baseline does not verify: %v", err)
			}

			// Flip every bit of the TCE bytes, one at a time.
			for i := 0; i < len(b); i++ {
				for _, mask := range []byte{0x01, 0x80} {
					mut := append([]byte{}, b...)
					mut[i] ^= mask
					if err := verifyFor(t, v.Name, reqTce, reqSig, rotTce, rotSig, mut, sig); err == nil {
						t.Fatalf("verified TCE with byte %d flipped (mask 0x%02x)", i, mask)
					}
				}
			}

			// Flip every bit of the signature, one at a time.
			for i := 0; i < len(sig); i++ {
				mut := append([]byte{}, sig...)
				mut[i] ^= 0x01
				if err := verifyFor(t, v.Name, reqTce, reqSig, rotTce, rotSig, b, mut); err == nil {
					t.Fatalf("verified with signature byte %d flipped", i)
				}
			}
		})
	}
}

func TestDenyAndAllowDifferInOneByte(t *testing.T) {
	vf := loadVectors(t)
	byName := map[string]*vectorEntry{}
	for i := range vf.Vectors {
		byName[vf.Vectors[i].Name] = &vf.Vectors[i]
	}
	allow := byName["approval_response/allow"]
	deny := byName["approval_response/deny"]
	req := byName["approval_request/ban"]

	a := mustHex(t, allow.TCEHex)
	d := mustHex(t, deny.TCEHex)
	if len(a) != len(d) {
		t.Fatalf("allow/deny differ in length")
	}
	diffs := 0
	for i := range a {
		if a[i] != d[i] {
			diffs++
		}
	}
	// The reference design makes deny differ from allow in exactly the
	// decision byte, so a verifier that ignores the decision is detectable.
	if diffs != 1 {
		t.Fatalf("expected the two responses to differ in one byte, got %d", diffs)
	}

	reqTce := mustHex(t, req.TCEHex)
	reqSig := mustHex(t, req.SignatureHex)
	allowSig := mustHex(t, allow.SignatureHex)
	denySig := mustHex(t, deny.SignatureHex)

	if _, err := protocol.VerifyApprovalResponse(reqTce, reqSig, d, denySig); err != nil {
		t.Fatalf("deny does not verify: %v", err)
	}
	if _, err := protocol.VerifyApprovalResponse(reqTce, reqSig, a, allowSig); err != nil {
		t.Fatalf("allow does not verify: %v", err)
	}
	// A decision cannot be transplanted across the two documents.
	if _, err := protocol.VerifyApprovalResponse(reqTce, reqSig, a, denySig); err == nil {
		t.Fatal("deny signature verified over allow bytes")
	}
	if _, err := protocol.VerifyApprovalResponse(reqTce, reqSig, d, allowSig); err == nil {
		t.Fatal("allow signature verified over deny bytes")
	}
}