package libsignal_test

// This is a self-consistency round trip, not a transcription of published
// X3DH/Double Ratchet test vectors: two simulated parties generate real
// keys through the cgo-linked library, establish a session via X3DH, and
// exchange messages, verifying the plaintext survives encrypt-on-one-side
// / decrypt-on-the-other for both protocol versions. That's the property
// that actually matters for this binding (it interoperates with the real
// library correctly), and avoids the risk of mistranscribing external
// vectors for a library/wire-format combination with no independently
// published vectors readily available.

import (
	"bytes"
	"testing"

	"dev.narayana.im/narayana/telegabber/e2ee/omemo/libsignal"
)

// identity bundles up one simulated party's generated keys and store.
type identity struct {
	ctx      *libsignal.Context
	storeCtx *libsignal.StoreContext

	registrationID     uint32
	identityPublicKey  []byte
	signedPreKeyID     uint32
	signedPreKeyRecord []byte
	preKeyID           uint32
	preKeyRecord       []byte
}

func newIdentity(t *testing.T) *identity {
	t.Helper()

	ctx, err := libsignal.NewContext()
	if err != nil {
		t.Fatalf("NewContext: %v", err)
	}
	t.Cleanup(ctx.Close)

	idKeyPair, err := libsignal.GenerateIdentityKeyPair(ctx)
	if err != nil {
		t.Fatalf("GenerateIdentityKeyPair: %v", err)
	}
	pub, priv, err := libsignal.SplitIdentityKeyPair(ctx, idKeyPair.Record)
	if err != nil {
		t.Fatalf("SplitIdentityKeyPair: %v", err)
	}
	regID, err := libsignal.GenerateRegistrationID(ctx)
	if err != nil {
		t.Fatalf("GenerateRegistrationID: %v", err)
	}

	store := newMemStore(pub, priv, regID)

	signedPreKey, err := libsignal.GenerateSignedPreKey(ctx, idKeyPair, 1)
	if err != nil {
		t.Fatalf("GenerateSignedPreKey: %v", err)
	}
	if err := store.StoreSignedPreKey(signedPreKey.ID, signedPreKey.Record); err != nil {
		t.Fatalf("StoreSignedPreKey: %v", err)
	}

	preKeys, err := libsignal.GeneratePreKeys(ctx, 1, 1)
	if err != nil {
		t.Fatalf("GeneratePreKeys: %v", err)
	}
	if len(preKeys) != 1 {
		t.Fatalf("expected 1 prekey, got %d", len(preKeys))
	}
	if err := store.StorePreKey(preKeys[0].ID, preKeys[0].Record); err != nil {
		t.Fatalf("StorePreKey: %v", err)
	}

	storeCtx, err := libsignal.NewStoreContext(ctx, store)
	if err != nil {
		t.Fatalf("NewStoreContext: %v", err)
	}
	t.Cleanup(storeCtx.Close)

	return &identity{
		ctx:                ctx,
		storeCtx:           storeCtx,
		registrationID:     regID,
		identityPublicKey:  pub,
		signedPreKeyID:     signedPreKey.ID,
		signedPreKeyRecord: signedPreKey.Record,
		preKeyID:           preKeys[0].ID,
		preKeyRecord:       preKeys[0].Record,
	}
}

// bundle returns id's own device bundle, as it would be published in an
// XEP-0384 bundle - i.e. what a remote peer's SessionBuilder needs to
// start a session with id via X3DH.
//
// Which signature to include depends on the recipient's protocol version:
// session_builder_process_pre_key_bundle re-serializes the signed prekey's
// public key via ec_public_key_serialize (version < 4, the 33-byte
// DJB_TYPE-prefixed legacy form) or ec_public_key_serialize_omemo
// (version >= 4, the raw 32-byte Montgomery form) before verifying it
// against whichever signature was supplied - confirmed by reading
// session_builder.c. Supplying the wrong one of the two signatures
// libomemo-c's signed-prekey generation always computes fails with
// SG_ERR_INVALID_KEY, not a version-related error, which is what actually
// surfaced this while writing this test.
func (id *identity) bundle(t *testing.T, deviceID uint32, version int) libsignal.RemoteBundle {
	t.Helper()

	spkInfo, err := libsignal.DecodeSignedPreKey(id.ctx, id.signedPreKeyRecord)
	if err != nil {
		t.Fatalf("DecodeSignedPreKey: %v", err)
	}
	pkInfo, err := libsignal.DecodePreKey(id.ctx, id.preKeyRecord)
	if err != nil {
		t.Fatalf("DecodePreKey: %v", err)
	}

	signature := spkInfo.Signature
	if version >= libsignal.ProtocolVersionModern {
		signature = spkInfo.SignatureOMEMO
	}

	return libsignal.RemoteBundle{
		RegistrationID:        id.registrationID,
		DeviceID:              deviceID,
		PreKeyID:              pkInfo.ID,
		PreKeyPublic:          pkInfo.PublicKey,
		SignedPreKeyID:        spkInfo.ID,
		SignedPreKeyPublic:    spkInfo.PublicKey,
		SignedPreKeySignature: signature,
		IdentityKeyPublic:     id.identityPublicKey,
	}
}

func TestSessionRoundTrip(t *testing.T) {
	for _, tc := range []struct {
		name    string
		version int
		omemo   bool
	}{
		{"legacy", libsignal.ProtocolVersionLegacy, false},
		{"modern", libsignal.ProtocolVersionModern, true},
	} {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			if tc.version == libsignal.ProtocolVersionModern && !libsignal.ModernOMEMOSupported {
				t.Skip("modern OMEMO not supported by this build (signal_legacy tag)")
			}

			alice := newIdentity(t)
			bob := newIdentity(t)

			aliceAddr := libsignal.Address{Name: "alice@example.com", DeviceID: 1}
			bobAddr := libsignal.Address{Name: "bob@example.com", DeviceID: 1}

			// Alice fetches Bob's bundle and establishes a session (X3DH).
			builder, err := libsignal.NewSessionBuilder(alice.ctx, alice.storeCtx, bobAddr, tc.version)
			if err != nil {
				t.Fatalf("NewSessionBuilder: %v", err)
			}
			defer builder.Close()
			if err := builder.ProcessPreKeyBundle(alice.ctx, bob.bundle(t, bobAddr.DeviceID, tc.version)); err != nil {
				t.Fatalf("ProcessPreKeyBundle: %v", err)
			}

			aliceCipher, err := libsignal.NewSessionCipher(alice.ctx, alice.storeCtx, bobAddr, tc.version)
			if err != nil {
				t.Fatalf("NewSessionCipher (alice): %v", err)
			}
			defer aliceCipher.Close()

			plaintext1 := []byte("hello bob")
			ct1, err := aliceCipher.Encrypt(plaintext1)
			if err != nil {
				t.Fatalf("Encrypt (alice->bob): %v", err)
			}
			if ct1.Type != libsignal.CiphertextPreKeyType {
				t.Fatalf("expected first message to be a prekey message, got type %d", ct1.Type)
			}

			// Bob decrypts Alice's first message - this establishes Bob's
			// side of the session as a side effect; no ProcessPreKeyBundle
			// needed on his end, only a bare SessionCipher.
			bobCipher, err := libsignal.NewSessionCipher(bob.ctx, bob.storeCtx, aliceAddr, tc.version)
			if err != nil {
				t.Fatalf("NewSessionCipher (bob): %v", err)
			}
			defer bobCipher.Close()

			got1, err := bobCipher.Decrypt(bob.ctx, *ct1, tc.omemo, bob.registrationID)
			if err != nil {
				t.Fatalf("Decrypt (bob<-alice): %v", err)
			}
			if !bytes.Equal(got1, plaintext1) {
				t.Fatalf("round trip mismatch: got %q, want %q", got1, plaintext1)
			}

			// Bob replies; Alice decrypts using her already-established
			// session (an ordinary ratchet message this time, not a
			// prekey message).
			plaintext2 := []byte("hi alice")
			ct2, err := bobCipher.Encrypt(plaintext2)
			if err != nil {
				t.Fatalf("Encrypt (bob->alice): %v", err)
			}
			if ct2.Type != libsignal.CiphertextSignalType {
				t.Fatalf("expected reply to be an ordinary ratchet message, got type %d", ct2.Type)
			}

			got2, err := aliceCipher.Decrypt(alice.ctx, *ct2, tc.omemo, alice.registrationID)
			if err != nil {
				t.Fatalf("Decrypt (alice<-bob): %v", err)
			}
			if !bytes.Equal(got2, plaintext2) {
				t.Fatalf("round trip mismatch: got %q, want %q", got2, plaintext2)
			}
		})
	}
}