Alice                                    Bob
  │                                        │
  │          Initial X3DH Setup            │
  │────────────────────────────────────────│
  │   Alice gets Bob’s Identity, PreKey,   │
  │   One-Time PreKey from server          │
  │   Computes 4 shared secrets → Root Key │
  │                                        │
  │◄───────────── Root Key ──────────────►│
  │                                        │
  ▼                                        ▼
Root Key → HKDF → Chain Key + Message Key   Chain Key + Message Key
       (Symmetric Ratchet starts)                (Symmetric Ratchet starts)

Alice sends message 1:
   • Uses current Message Key → encrypt
   • Then: Chain Key = HKDF(Chain Key, "ratchet")
   • New Message Key = HKDF(Chain Key, "message")

Bob receives → decrypts → updates his Chain Key same way

Alice sends message 2 → same symmetric ratchet

Now Bob replies → NEW DH Ratchet triggers!
   • Bob sends his new Ratchet Public Key
   • Both do new DH with their private + other's new public
   • New Root Key = HKDF(old Root Key + new DH)
   • Symmetric chains RESET with new Chain Key
   • Old Chain Keys DELETED forever → Forward Secrecy achieved

Every time someone replies → new DH ratchet → old keys die
Even if no reply → symmetric ratchet keeps burning old keys

# double_ratchet_mini.py  ← Run this in lab → impress everyone
from cryptography.hazmat.primitives.asymmetric import x25519
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.fernet import Fernet
import os

class DoubleRatchet:
    def __init__(self, initial_root_key):
        self.root_key = initial_root_key
        self.send_chain_key = None
        self.recv_chain_key = None
        self.dh_private = x25519.X25519PrivateKey.generate()
        self.dh_public = self.dh_private.public_key()
        self.peer_dh_public = None
        self.skipped_mk = {}  # for out-of-order messages

    def dh_ratchet_step(self, peer_public_key_bytes):
        peer_pub = x25519.X25519PublicKey.from_public_bytes(peer_public_key_bytes)
        shared = self.dh_private.exchange(peer_pub)

        # Root key update
        hkdf = HKDF(hashes.SHA256(), 64, salt=self.root_key, info=b"ratchet")
        keys = hkdf.derive(shared)
        self.root_key = keys[:32]
        self.recv_chain_key = keys[32:]

        # Prepare to send next message
        self.dh_private = x25519.X25519.X25519PrivateKey.generate()
        self.dh_public = self.dh_private.public_key()

    def send_message(self, plaintext):
        if self.send_chain_key is None:
            self.send_chain_key = self.root_key  # first time

        hkdf = HKDF(hashes.SHA256(), 32, salt=None, info=b"msgkey")
        msg_key = hkdf.derive(self.send_chain_key + b"send")
        cipher = Fernet(msg_key).encrypt(plaintext.encode())

        # Advance sending chain
        hkdf_chain = HKDF(hashes.SHA256(), 32, salt=None, info=b"chain")
        self.send_chain_key = hkdf_chain.derive(self.send_chain_key)

        return cipher, self.dh_public.public_bytes_raw()

    def receive_message(self, ciphertext, peer_dh_public_bytes):
        # DH Ratchet if new public key
        if self.peer_dh_public is None or peer_dh_public_bytes != self.peer_dh_public:
            self.dh_ratchet_step(peer_dh_public_bytes)
            self.peer_dh_public = peer_dh_public_bytes

        # Use receiving chain key
        hkdf = HKDF(hashes.SHA256(), 32, salt=None, info=b"msgkey")
        msg_key = hkdf.derive(self.recv_chain_key + b"recv")

        try:
            plain = Fernet(msg_key).decrypt(ciphertext).decode()
            # Advance receiving chain
            hkdf_chain = HKDF(hashes.SHA256(), 32, salt=None, info=b"chain")
            self.recv_chain_key = hkdf_chain.derive(self.recv_chain_key)
            return plain
        except:
            return "[Decryption failed]"

# Demo
import os
initial_shared = os.urandom(32)  # from X3DH

alice = DoubleRatchet(initial_shared)
bob = DoubleRatchet(initial_shared)

# Alice sends first message
c1, alice_pub1 = alice.send_message("Hello Bob, this is secret!")
print("Alice → Bob:", bob.receive_message(c1, alice_pub1))

# Bob replies → triggers DH ratchet
c2, bob_pub = bob.send_message("Hi Alice, I love Double Ratchet!")
print("Bob → Alice:", alice.receive_message(c2, bob_pub))

# Alice sends another
c3, alice_pub2 = alice.send_message("Even if you steal my phone now, past messages are safe!")
print("Alice → Bob:", bob.receive_message(c3, alice_pub2))