web/backend/services/messaging/processor.py

"""
Message processing pipeline for envelope encryption.

This module handles the core envelope encryption workflow:
1. Decrypt client-encrypted message (transport encryption)
2. Generate random MEK
3. Encrypt plaintext with MEK
4. Wrap MEK for compliance, sender, and recipient
5. Store encrypted message + wrapped keys
"""

import logging
import json
import time
import base64
from typing import Dict, Any, Optional
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey

from .encryption import (
    decrypt_transport_message,
    generate_mek,
    encrypt_message,
    wrap_mek,
    derive_shared_secret,
    derive_key_from_shared_secret,
)

logger = logging.getLogger("uvicorn.error")


def process_encrypted_message(
    client_public_key_b64: str,
    transport_nonce_b64: str,
    transport_ciphertext_b64: str,
    compliance_public_key_b64: str,
    sender_public_key_b64: str,
    recipient_public_key_b64: str,
    ephemeral_private_key: X25519PrivateKey,
) -> Dict[str, Any]:
    """
    Process an encrypted message through the envelope encryption pipeline.
    
    Step 1: Decrypt client message using transport encryption (ephemeral keys)
    Step 2: Generate random MEK
    Step 3: Encrypt plaintext with MEK
    Step 4: Wrap MEK for compliance, sender, recipient (using their provided public keys)
    Step 5: Return encrypted message + 3 wrapped MEKs
    
    Args:
        client_public_key_b64: Client's ephemeral public key for transport decryption
        transport_nonce_b64: Nonce used for transport encryption
        transport_ciphertext_b64: Client's encrypted plaintext
        compliance_public_key_b64: Compliance system's public key for MEK wrapping
        sender_public_key_b64: Sender's public key for MEK wrapping
        recipient_public_key_b64: Recipient's public key for MEK wrapping
        ephemeral_private_key: Server's ephemeral X25519 private key
    
    Returns:
        Dict with encrypted message and wrapped MEKs:
        {
            "nonce": base64-encoded nonce for content encryption,
            "ciphertext": base64-encoded encrypted content,
            "compliance_wrapped_mek": base64-encoded wrapped MEK,
            "sender_wrapped_mek": base64-encoded wrapped MEK,
            "recipient_wrapped_mek": base64-encoded wrapped MEK,
        }
    """
    try:
        start_time = time.time()
        
        # Step 1: Decrypt transport message
        logger.info("CRYPTO: Starting envelope encryption processing")
        plaintext = decrypt_transport_message(
            client_public_key_b64,
            transport_nonce_b64,
            transport_ciphertext_b64,
            ephemeral_private_key,
        )
        logger.info(
            "CRYPTO: Transport decryption complete, plaintext size: %d bytes",
            len(plaintext)
        )

        # Step 2: Generate random MEK
        mek = generate_mek()
        logger.info("CRYPTO: Generated random MEK (32 bytes)")

        # Step 3: Encrypt plaintext with MEK
        content_nonce, ciphertext = encrypt_message(plaintext, mek)
        logger.info(
            "CRYPTO: Content encryption with MEK complete, ciphertext size: %d bytes",
            len(ciphertext)
        )

        # Step 4a: Derive wrap keys deterministically from recipient public keys
        # This avoids needing to store the ephemeral transport key
        logger.info("CRYPTO: Deriving key wrap keys deterministically")

        # Use HKDF with recipient public key bytes as input to derive wrap keys
        # This is deterministic and doesn't require storing ephemeral keys
        import base64
        compliance_key_bytes = base64.b64decode(compliance_public_key_b64)
        sender_key_bytes = base64.b64decode(sender_public_key_b64)
        recipient_key_bytes = base64.b64decode(recipient_public_key_b64)

        logger.info(f"🔑 Deriving wrap keys for sender={sender_public_key_b64[:20]}... recipient={recipient_public_key_b64[:20]}...")

        compliance_wrap_key = derive_key_from_shared_secret(compliance_key_bytes, "compliance_wrap_key")
        sender_wrap_key = derive_key_from_shared_secret(sender_key_bytes, "sender_wrap_key")
        recipient_wrap_key = derive_key_from_shared_secret(recipient_key_bytes, "recipient_wrap_key")

        logger.info("✅ Wrap keys derived successfully")

        # Step 4b: Wrap MEK for each recipient
        compliance_wrapped_mek = wrap_mek(mek, compliance_wrap_key)
        sender_wrapped_mek = wrap_mek(mek, sender_wrap_key)
        recipient_wrapped_mek = wrap_mek(mek, recipient_wrap_key)

        duration = time.time() - start_time
        logger.info(
            "CRYPTO: Successfully processed message with 3 MEK wraps (compliance/sender/recipient) in %.2fms",
            duration * 1000
        )

        # Get the transport public key for storage with the message
        transport_public_key_b64 = base64.b64encode(ephemeral_private_key.public_key().public_bytes_raw()).decode("ascii")

        return {
            "nonce": content_nonce,
            "ciphertext": ciphertext,
            "compliance_wrapped_mek": compliance_wrapped_mek,
            "sender_wrapped_mek": sender_wrapped_mek,
            "recipient_wrapped_mek": recipient_wrapped_mek,
        }

    except Exception as e:
        duration = time.time() - start_time
        logger.exception(
            "CRYPTO: Failed to process encrypted message after %.2fms: %s",
            duration * 1000, str(e)
        )
        raise


def process_encrypted_message_and_files(
    plaintext_message: bytes,
    plaintext_files: list[bytes],
    compliance_public_key_b64: str,
    sender_public_key_b64: str,
    recipient_public_key_b64: str,
) -> Dict[str, Any]:
    """
    Process a message and its attached files using a single MEK.

    - Generates one random MEK
    - Encrypts message and each file with AES-GCM using that MEK (unique nonce per item)
    - Wraps the MEK for compliance, sender, and recipient

    Returns:
        {
            "message": {"nonce": str, "ciphertext": str},
            "files": [{"nonce": str, "ciphertext": str}, ...],
            "compliance_wrapped_mek": str,
            "sender_wrapped_mek": str,
            "recipient_wrapped_mek": str,
        }
    """
    start_time = time.time()

    # One MEK for everything in this envelope
    mek = generate_mek()

    # Encrypt message
    msg_nonce, msg_ciphertext = encrypt_message(plaintext_message, mek)

    # Encrypt files (same MEK, per-file nonce)
    files_out: list[Dict[str, str]] = []
    for f_bytes in plaintext_files:
        f_nonce, f_ciphertext = encrypt_message(f_bytes, mek)
        files_out.append({"nonce": f_nonce, "ciphertext": f_ciphertext})

    # Derive wrap keys deterministically (same as existing flow)
    compliance_key_bytes = base64.b64decode(compliance_public_key_b64)
    sender_key_bytes = base64.b64decode(sender_public_key_b64)
    recipient_key_bytes = base64.b64decode(recipient_public_key_b64)

    compliance_wrap_key = derive_key_from_shared_secret(compliance_key_bytes, "compliance_wrap_key")
    sender_wrap_key = derive_key_from_shared_secret(sender_key_bytes, "sender_wrap_key")
    recipient_wrap_key = derive_key_from_shared_secret(recipient_key_bytes, "recipient_wrap_key")

    compliance_wrapped_mek = wrap_mek(mek, compliance_wrap_key)
    sender_wrapped_mek = wrap_mek(mek, sender_wrap_key)
    recipient_wrapped_mek = wrap_mek(mek, recipient_wrap_key)

    duration = time.time() - start_time
    logger.info(
        "CRYPTO: Processed message+%d files with single MEK in %.2fms",
        len(files_out),
        duration * 1000,
    )

    return {
        "message": {"nonce": msg_nonce, "ciphertext": msg_ciphertext},
        "files": files_out,
        "compliance_wrapped_mek": compliance_wrapped_mek,
        "sender_wrapped_mek": sender_wrapped_mek,
        "recipient_wrapped_mek": recipient_wrapped_mek,
    }