featherChat/warzone/crates/warzone-client/src/cli/recv.rs

use anyhow::{Context, Result};
use warzone_protocol::identity::IdentityKeyPair;
use warzone_protocol::ratchet::RatchetState;
use warzone_protocol::types::Fingerprint;
use warzone_protocol::x3dh;
use x25519_dalek::PublicKey;

use crate::cli::send::WireMessage;
use crate::net::ServerClient;
use crate::storage::LocalDb;

pub async fn run(server_url: &str, identity: &IdentityKeyPair) -> Result<()> {
    let our_pub = identity.public_identity();
    let our_fp = our_pub.fingerprint.to_string();
    let db = LocalDb::open()?;
    let client = ServerClient::new(server_url);

    println!("Polling for messages as {}...", our_fp);

    let messages = client.poll_messages(&our_fp).await?;

    if messages.is_empty() {
        println!("No new messages.");
        return Ok(());
    }

    println!("Received {} message(s):\n", messages.len());

    for raw in &messages {
        match bincode::deserialize::<WireMessage>(raw) {
            Ok(WireMessage::KeyExchange {
                sender_fingerprint,
                sender_identity_encryption_key,
                ephemeral_public,
                used_one_time_pre_key_id,
                ratchet_message,
            }) => {
                let sender_fp = Fingerprint::from_hex(&sender_fingerprint)
                    .context("invalid sender fingerprint")?;

                // Load our signed pre-key secret
                let spk_id = 1u32; // We use ID 1 for our signed pre-key
                let spk_secret = db
                    .load_signed_pre_key(spk_id)?
                    .context("missing signed pre-key — run `warzone init` first")?;

                // Load one-time pre-key if used
                let otpk_secret = if let Some(id) = used_one_time_pre_key_id {
                    db.take_one_time_pre_key(id)?
                } else {
                    None
                };

                // X3DH respond
                let their_identity_x25519 = PublicKey::from(sender_identity_encryption_key);
                let their_ephemeral = PublicKey::from(ephemeral_public);

                let shared_secret = x3dh::respond(
                    &identity,
                    &spk_secret,
                    otpk_secret.as_ref(),
                    &their_identity_x25519,
                    &their_ephemeral,
                )
                .context("X3DH respond failed")?;

                // Init ratchet as Bob
                let mut state = RatchetState::init_bob(shared_secret, spk_secret);

                // Decrypt the message
                match state.decrypt(&ratchet_message) {
                    Ok(plaintext) => {
                        let text = String::from_utf8_lossy(&plaintext);
                        println!("  [{}] {}: {}", "new session", sender_fingerprint, text);
                        db.save_session(&sender_fp, &state)?;
                    }
                    Err(e) => {
                        eprintln!("  [{}] decrypt failed: {}", sender_fingerprint, e);
                    }
                }
            }
            Ok(WireMessage::Message {
                sender_fingerprint,
                ratchet_message,
            }) => {
                let sender_fp = Fingerprint::from_hex(&sender_fingerprint)
                    .context("invalid sender fingerprint")?;

                match db.load_session(&sender_fp)? {
                    Some(mut state) => match state.decrypt(&ratchet_message) {
                        Ok(plaintext) => {
                            let text = String::from_utf8_lossy(&plaintext);
                            println!("  {}: {}", sender_fingerprint, text);
                            db.save_session(&sender_fp, &state)?;
                        }
                        Err(e) => {
                            eprintln!("  [{}] decrypt failed: {}", sender_fingerprint, e);
                        }
                    },
                    None => {
                        eprintln!(
                            "  [{}] no session — cannot decrypt (need key exchange first)",
                            sender_fingerprint
                        );
                    }
                }
            }
            Err(e) => {
                eprintln!("  failed to deserialize message: {}", e);
            }
        }
    }

    Ok(())
}