wacore 0.7.0

//! E2E-encrypted message edit envelope (`secret_encrypted_message` with
//! `secret_enc_type = MESSAGE_EDIT`).
//!
//! Introduced by WhatsApp in 2026: newer clients wrap message edits in an
//! E2EE envelope keyed by the original message's
//! `messageContextInfo.messageSecret`, replacing the older in-clear
//! `protocolMessage.editedMessage` path.
//!
//! Verified against `docs/captured-js/`:
//!
//! - `WA/Use/CaseSecret.js` — use-case literal `"Message Edit"` and the
//!   HKDF info ordering (`stanzaId || parentOrigSender || editor || usecase`).
//! - `WAWeb/Parse/MessageEditEncryptedMessageProto.js` — envelope detection
//!   and IV-length validation (`u.length !== 12`).
//! - `WAWeb/Process/EncryptedMessageEditMsgs.js` — invocation contract.
//! - `WAWeb/Addon/Encryption.js` function `g` — AAD is empty for the
//!   MessageEdit branch (only `PollVote` and `EventResponse` bind stanza/sender
//!   into AAD).
//!
//! The plaintext is a `Message` proto whose `protocolMessage.editedMessage`
//! carries the new content — same shape as the legacy edit path, so callers
//! that already handle `protocolMessage.editedMessage` can reuse their code.

use anyhow::{Result, anyhow};

use crate::secret_enc_addon::{AddonContext, ModificationType, decrypt_addon, encrypt_addon};

const IV_SIZE: usize = 12;

/// Inputs for decrypting / encrypting a `MESSAGE_EDIT` envelope.
///
/// JIDs must be passed in the addressing they were received under. The
/// caller is responsible for any LID↔PN normalisation; see
/// [`decrypt_message_edit_with_fallback`].
#[derive(Debug, Clone, Copy)]
pub struct MessageEditContext<'a> {
    /// ID of the *original* (target) message being edited — this is the
    /// stanza id WA Web feeds into the HKDF info buffer.
    pub original_msg_id: &'a str,
    /// JID of the sender of the original message (`participant` for groups,
    /// `remote_jid` for 1:1).
    pub original_sender_jid: &'a str,
    /// JID of the user performing the edit (sender of the envelope).
    pub editor_jid: &'a str,
}

impl<'a> MessageEditContext<'a> {
    fn as_addon_ctx(&self) -> AddonContext<'a> {
        AddonContext {
            stanza_id: self.original_msg_id,
            parent_msg_original_sender: self.original_sender_jid,
            modification_sender: self.editor_jid,
            modification_type: ModificationType::MessageEdit,
        }
    }
}

/// Encrypt an edit payload for testing or for client-side outgoing edits.
///
/// `inner_message` is the full `Message` proto whose `protocolMessage.editedMessage`
/// carries the new content; it gets serialised and encrypted in one shot.
pub fn encrypt_message_edit(
    inner_message: &waproto::whatsapp::Message,
    message_secret: &[u8],
    ctx: &MessageEditContext<'_>,
) -> Result<(Vec<u8>, [u8; IV_SIZE])> {
    let plaintext = waproto::codec::message_to_vec(inner_message);
    encrypt_addon(&plaintext, message_secret, &ctx.as_addon_ctx())
}

/// Decrypt any `secret_encrypted_message` envelope to its inner `Message`.
///
/// All `SecretEncType` envelope variants (`EVENT_EDIT`, `MESSAGE_EDIT`,
/// `POLL_EDIT`, `POLL_ADD_OPTION`) share the same shape: they decrypt to a full
/// `Message` proto with an empty AAD (per `WAWebAddonEncryption`'s function `g`),
/// differing only in the use-case secret fed into the HKDF — which is exactly
/// what `modification_type` selects. Mirrors whatsmeow's
/// `DecryptSecretEncryptedMessage` dispatch.
///
/// `iv` must be exactly 12 bytes (matches WA Web's `u.length !== 12` check).
pub fn decrypt_secret_encrypted(
    enc_payload: &[u8],
    iv: &[u8],
    message_secret: &[u8],
    modification_type: ModificationType,
    ctx: &MessageEditContext<'_>,
) -> Result<waproto::whatsapp::Message> {
    if iv.len() != IV_SIZE {
        return Err(anyhow!(
            "Invalid secret-encrypted IV length: expected {IV_SIZE}, got {}",
            iv.len()
        ));
    }
    let addon = AddonContext {
        stanza_id: ctx.original_msg_id,
        parent_msg_original_sender: ctx.original_sender_jid,
        modification_sender: ctx.editor_jid,
        modification_type,
    };
    let plaintext = decrypt_addon(enc_payload, iv, message_secret, &addon)?;
    waproto::codec::message_decode(&plaintext[..])
        .map_err(|e| anyhow!("Failed to decode inner secret-encrypted Message: {e}"))
}

/// Decrypt a `secret_encrypted_message` MESSAGE_EDIT envelope.
///
/// Returns the inner `Message` whose `protocolMessage.editedMessage` carries
/// the new content. The caller can re-wrap it in `protocolMessage.editedMessage`
/// for consumer parity with the legacy edit path.
///
/// `iv` must be exactly 12 bytes (matches WA Web's `u.length !== 12` check).
pub fn decrypt_message_edit(
    enc_payload: &[u8],
    iv: &[u8],
    message_secret: &[u8],
    ctx: &MessageEditContext<'_>,
) -> Result<waproto::whatsapp::Message> {
    decrypt_secret_encrypted(
        enc_payload,
        iv,
        message_secret,
        ModificationType::MessageEdit,
        ctx,
    )
}

/// Decrypt a secret-encrypted envelope, retrying once under an alternate
/// addressing, mirroring WA Web's `decryptAddOn` LID/PN resilience.
///
/// `primary` is the as-received JID pair; `fallback` is the swapped pair.
/// Returns the first success. WA Web tries 3 attempts (LID→LID, PN→PN,
/// originals) but two suffice when the caller already knows the canonical pair.
pub fn decrypt_secret_encrypted_with_fallback(
    enc_payload: &[u8],
    iv: &[u8],
    message_secret: &[u8],
    modification_type: ModificationType,
    primary: &MessageEditContext<'_>,
    fallback: Option<&MessageEditContext<'_>>,
) -> Result<waproto::whatsapp::Message> {
    match decrypt_secret_encrypted(enc_payload, iv, message_secret, modification_type, primary) {
        Ok(m) => Ok(m),
        Err(primary_err) => match fallback {
            Some(fb) => decrypt_secret_encrypted(
                enc_payload,
                iv,
                message_secret,
                modification_type,
                fb,
            )
            .map_err(|fb_err| {
                anyhow!("secret-encrypted decrypt failed: primary={primary_err}; fallback={fb_err}")
            }),
            None => Err(primary_err),
        },
    }
}

/// MESSAGE_EDIT-specialised wrapper over
/// [`decrypt_secret_encrypted_with_fallback`].
pub fn decrypt_message_edit_with_fallback(
    enc_payload: &[u8],
    iv: &[u8],
    message_secret: &[u8],
    primary: &MessageEditContext<'_>,
    fallback: Option<&MessageEditContext<'_>>,
) -> Result<waproto::whatsapp::Message> {
    decrypt_secret_encrypted_with_fallback(
        enc_payload,
        iv,
        message_secret,
        ModificationType::MessageEdit,
        primary,
        fallback,
    )
}

#[cfg(test)]
#[allow(clippy::disallowed_methods)]
mod tests {
    use super::*;
    use buffa::MessageField;
    use waproto::whatsapp as wa;

    fn make_inner_edit(new_text: &str) -> wa::Message {
        wa::Message {
            protocol_message: MessageField::some(wa::message::ProtocolMessage {
                key: MessageField::some(wa::MessageKey {
                    remote_jid: Some("g@g.us".to_string()),
                    from_me: Some(true),
                    id: Some("AC1234567890ABCDEF".to_string()),
                    participant: None,
                }),
                r#type: Some(wa::message::protocol_message::Type::MESSAGE_EDIT),
                edited_message: MessageField::some(wa::Message {
                    conversation: Some(new_text.to_string()),
                    ..Default::default()
                }),
                timestamp_ms: Some(1_700_000_000_000),
                ..Default::default()
            }),
            ..Default::default()
        }
    }

    #[test]
    fn encrypt_decrypt_roundtrip_text() {
        let secret = [0x07u8; 32];
        let ctx = MessageEditContext {
            original_msg_id: "AC1234567890ABCDEF",
            original_sender_jid: "5511999999999@s.whatsapp.net",
            editor_jid: "5511999999999@s.whatsapp.net",
        };
        let inner = make_inner_edit("edited text");
        let (enc, iv) = encrypt_message_edit(&inner, &secret, &ctx).unwrap();

        let decoded = decrypt_message_edit(&enc, &iv, &secret, &ctx).unwrap();
        let edited = decoded
            .protocol_message
            .as_option()
            .and_then(|pm| pm.edited_message.as_option())
            .expect("inner edited message present");
        assert_eq!(edited.conversation.as_deref(), Some("edited text"));
    }

    #[test]
    fn encrypt_decrypt_roundtrip_with_lid_jids() {
        let secret = [0x42u8; 32];
        let ctx = MessageEditContext {
            original_msg_id: "AC1191FE0A25A0E319BEA72064819280",
            original_sender_jid: "260661598801930@lid",
            editor_jid: "260661598801930@lid",
        };
        let inner = make_inner_edit("B");
        let (enc, iv) = encrypt_message_edit(&inner, &secret, &ctx).unwrap();
        let decoded = decrypt_message_edit(&enc, &iv, &secret, &ctx).unwrap();
        assert_eq!(
            decoded
                .protocol_message
                .as_option()
                .and_then(|pm| pm.edited_message.as_option())
                .and_then(|m| m.conversation.as_deref()),
            Some("B")
        );
    }

    #[test]
    fn wrong_editor_jid_fails() {
        // GCM tag should fail when the editor JID feeding HKDF differs.
        let secret = [0x07u8; 32];
        let ctx = MessageEditContext {
            original_msg_id: "AC1",
            original_sender_jid: "a@s.whatsapp.net",
            editor_jid: "a@s.whatsapp.net",
        };
        let (enc, iv) = encrypt_message_edit(&make_inner_edit("x"), &secret, &ctx).unwrap();

        let bad = MessageEditContext {
            editor_jid: "b@s.whatsapp.net",
            ..ctx
        };
        assert!(decrypt_message_edit(&enc, &iv, &secret, &bad).is_err());
    }

    #[test]
    fn wrong_message_secret_fails() {
        let ctx = MessageEditContext {
            original_msg_id: "AC1",
            original_sender_jid: "a@s.whatsapp.net",
            editor_jid: "a@s.whatsapp.net",
        };
        let (enc, iv) = encrypt_message_edit(&make_inner_edit("x"), &[0x07u8; 32], &ctx).unwrap();
        assert!(decrypt_message_edit(&enc, &iv, &[0x08u8; 32], &ctx).is_err());
    }

    #[test]
    fn invalid_iv_length_rejected() {
        let ctx = MessageEditContext {
            original_msg_id: "AC1",
            original_sender_jid: "a@s.whatsapp.net",
            editor_jid: "a@s.whatsapp.net",
        };
        let (enc, _iv) = encrypt_message_edit(&make_inner_edit("x"), &[0x07u8; 32], &ctx).unwrap();
        // WA Web enforces 12-byte IV; we surface a typed error here.
        assert!(decrypt_message_edit(&enc, &[0u8; 11], &[0x07u8; 32], &ctx).is_err());
        assert!(decrypt_message_edit(&enc, &[0u8; 16], &[0x07u8; 32], &ctx).is_err());
    }

    #[test]
    fn general_decrypt_roundtrips_non_edit_use_case() {
        use crate::secret_enc_addon::{AddonContext, ModificationType, encrypt_addon};
        use buffa::Message as _;

        // A POLL_EDIT envelope: same shape as MESSAGE_EDIT, different use-case.
        let secret = [0x71u8; 32];
        let ctx = MessageEditContext {
            original_msg_id: "POLLID",
            original_sender_jid: "creator@s.whatsapp.net",
            editor_jid: "editor@s.whatsapp.net",
        };
        let inner = wa::Message {
            conversation: Some("poll edited".to_string()),
            ..Default::default()
        };
        let (enc, iv) = encrypt_addon(
            &inner.encode_to_vec(),
            &secret,
            &AddonContext {
                stanza_id: ctx.original_msg_id,
                parent_msg_original_sender: ctx.original_sender_jid,
                modification_sender: ctx.editor_jid,
                modification_type: ModificationType::PollEdit,
            },
        )
        .unwrap();

        // Correct use-case decrypts; the MESSAGE_EDIT wrapper (wrong use-case) fails.
        let out =
            decrypt_secret_encrypted(&enc, &iv, &secret, ModificationType::PollEdit, &ctx).unwrap();
        assert_eq!(out.conversation.as_deref(), Some("poll edited"));
        assert!(decrypt_message_edit(&enc, &iv, &secret, &ctx).is_err());
    }

    #[test]
    fn fallback_recovers_on_alternate_jid_form() {
        let secret = [0x09u8; 32];
        // Encrypted under PN form...
        let pn_ctx = MessageEditContext {
            original_msg_id: "ID",
            original_sender_jid: "5511999@s.whatsapp.net",
            editor_jid: "5511999@s.whatsapp.net",
        };
        // ...but consumer first guesses LID.
        let lid_ctx = MessageEditContext {
            original_msg_id: "ID",
            original_sender_jid: "12345@lid",
            editor_jid: "12345@lid",
        };

        let (enc, iv) = encrypt_message_edit(&make_inner_edit("hello"), &secret, &pn_ctx).unwrap();

        // Primary (LID) fails, fallback (PN) succeeds.
        let m = decrypt_message_edit_with_fallback(&enc, &iv, &secret, &lid_ctx, Some(&pn_ctx))
            .expect("fallback should rescue");
        assert_eq!(
            m.protocol_message
                .as_option()
                .and_then(|pm| pm.edited_message.as_option())
                .and_then(|m| m.conversation.as_deref()),
            Some("hello")
        );
    }

    #[test]
    fn fallback_returns_combined_error_when_both_fail() {
        let secret = [0x09u8; 32];
        let pn_ctx = MessageEditContext {
            original_msg_id: "ID",
            original_sender_jid: "5511999@s.whatsapp.net",
            editor_jid: "5511999@s.whatsapp.net",
        };
        let (enc, iv) = encrypt_message_edit(&make_inner_edit("x"), &secret, &pn_ctx).unwrap();

        let wrong1 = MessageEditContext {
            editor_jid: "evil1@s.whatsapp.net",
            ..pn_ctx
        };
        let wrong2 = MessageEditContext {
            editor_jid: "evil2@s.whatsapp.net",
            ..pn_ctx
        };
        let err = decrypt_message_edit_with_fallback(&enc, &iv, &secret, &wrong1, Some(&wrong2))
            .expect_err("both fail");
        let s = err.to_string();
        assert!(s.contains("primary="));
        assert!(s.contains("fallback="));
    }
}