Ironforce/src/message.rs

use alloc::string::String;
use alloc::vec::Vec;
use crate::crypto::{Keys, PublicKey};
use crate::res::IFResult;
use crate::tunnel::TunnelPublic;
use serde::{Serialize, Deserialize};
use sha2::Digest;


/// A serialized message
pub(crate) type MessageBytes = Vec<u8>;

/// Signature of the message: optional and optionally encrypted sender's key and signed hash
#[derive(Serialize, Deserialize, Clone)]
pub enum Signature {
    /// The message is signed. Author is unknown
    NotSigned,
    /// The message is signed with the sender's key visible to everyone
    Signed {
        sender: PublicKey,
        signature: Vec<u8>,
    },
    /// Sender's key is encrypted for the recipient
    SignedPrivately {
        sender_encrypted: Vec<u8>,
        signature: Vec<u8>,
    },
}

/// Network name and version
#[derive(Serialize, Deserialize, Clone)]
pub struct NetworkInfo {
    network_name: String,
    version: String,
}

impl Default for NetworkInfo {
    fn default() -> Self {
        Self { version: String::from("0.1.0"), network_name: String::from("test") }
    }
}

#[derive(Serialize, Deserialize, Clone)]
pub enum MessageType {
    SingleCast,
    Broadcast,
    Service(ServiceMessageType),
}

impl MessageType {
    fn hash(&self) -> Vec<u8> {
        match self {
            MessageType::SingleCast => Vec::from([0]),
            MessageType::Broadcast => Vec::from([1]),
            MessageType::Service(ServiceMessageType::TunnelBuilding(tunnel)) => [2, 0].iter().chain(tunnel.hash().iter()).copied().collect()
        }
    }
}

#[derive(Serialize, Deserialize, Clone)]
pub enum ServiceMessageType {
    TunnelBuilding(TunnelPublic)
}

#[derive(Serialize, Deserialize, Clone)]
pub enum MessageContent {
    /// Just plaintext message content
    Plain(Vec<u8>),
    /// Message content bytes encrypted for the recipient
    Encrypted(Vec<u8>),
    /// There is no content
    None,
}

impl MessageContent {
    pub fn hash(&self) -> Vec<u8> {
        match self {
            MessageContent::Plain(v) => sha2::Sha512::new()
                .chain(&[0u8; 1])
                .chain(v.as_slice())
                .result().to_vec(),
            MessageContent::Encrypted(v) => sha2::Sha512::new()
                .chain(&[1; 1])
                .chain(v.as_slice())
                .result().to_vec(),
            MessageContent::None => Vec::new()
        }
    }
}

/// The struct for messages that are sent in the network
#[derive(Serialize, Deserialize, Clone)]
pub struct Message {
    /// Content of the message (not to be confused with the bytes that we are sending through interfaces)
    ///
    /// AKA useful payload
    pub content: MessageContent,
    /// The type of this message
    pub message_type: MessageType,
    /// Sender's signature
    pub signature: Signature,
    /// A random number that is used in hash together with the content
    salt: u64,
    /// Hash of message content and the salt
    hash: Vec<u8>,
    /// Optional: hash of the message encrypted for the recipient, so that the recipient can know that this message is for them, but nobody else
    recipient_verification: Option<Vec<u8>>,
    /// ID of the tunnel that is used
    tunnel_id: u64,
    /// Network info
    network_info: NetworkInfo,
}

impl Message {
    /// Verify message's hash
    pub fn verify(&self) -> bool {
        todo!()
    }

    /// Check if this message is for this set of keys
    pub fn check_recipient(&self, _keys: Keys) -> bool {
        todo!()
    }

    /// Get decrypted content of the message
    pub fn get_decrypted(&self, _keys: Keys) -> IFResult<Vec<u8>> {
        todo!()
    }

    pub fn calculate_hash(content: &MessageContent, message_type: MessageType, sender_or_encrypted_sender: Option<Vec<u8>>, network_info: &NetworkInfo) -> Vec<u8> {
        sha2::Sha512::new()
            .chain(content.hash().as_slice())
            .chain(message_type.hash().as_slice())
            .chain(sender_or_encrypted_sender.unwrap_or_default().as_slice())
            .chain(network_info.network_name.as_bytes())
            .chain(network_info.version.as_bytes())
            .result().to_vec()
    }

    /// Encrypt hash of the message for the recipient
    pub fn generate_recipient_verification(hash: Vec<u8>, recipient: PublicKey) -> rsa::errors::Result<Vec<u8>> {
        recipient.encrypt_data(&hash)
    }
}


/// Message builder to create a new message step-by-step, like `Message::build().message_type(...).sign(...)`
pub struct MessageBuilder {
    content: MessageContent,
    /// The type of the message to be built
    message_type: Option<MessageType>,
    /// Sender's keys
    sender: Option<Keys>,
    /// Recipient's public key (if present, the content will be encrypted and recipient verification field will be set)
    recipient: Option<PublicKey>,
    /// ID of the tunnel that is used
    tunnel_id: u64,
}

impl MessageBuilder {
    /// Create a new `MessageBuilder` with default parameters
    pub fn new() -> Self {
        Self {
            content: MessageContent::None,
            message_type: None,
            sender: None,
            recipient: None,
            tunnel_id: 0,
        }
    }

    pub fn content(mut self, cont: Vec<u8>) -> Self {
        self.content = MessageContent::Plain(cont);
        self
    }

    /// Sign the message
    pub fn sign(mut self, keys: &Keys) -> Self {
        self.sender = Some(keys.clone());
        self
    }

    /// Set message's recipient (and therefore set recipient verification and encrypt the content)
    pub fn recipient(mut self, recipient: PublicKey) -> Self {
        self.recipient = Some(recipient);
        self
    }

    /// Set tunnel id
    pub fn tunnel(mut self, tunnel_id: u64) -> Self {
        self.tunnel_id = tunnel_id;
        self
    }

    /// Set message's type
    pub fn message_type(mut self, message_type: MessageType) -> Self {
        self.message_type = Some(message_type);
        self
    }

    /// Get the resulting message
    pub fn build(self) -> IFResult<Message> {
        let salt = rand::random();
        let sender_encrypted = if let (Some(sender_keys), Some(recipient)) = (self.sender.as_ref(), self.recipient.as_ref()) {
            Some(recipient.encrypt_data(&sender_keys.get_public().to_vec()?)?)
        } else { None };
        let network_info = NetworkInfo::default();
        let hash = Message::calculate_hash(
            &self.content,
            self.message_type.clone().unwrap(),
            sender_encrypted.clone()
                .or_else(
                    || self.sender.as_ref()
                        .map(|sender_keys| sender_keys.get_public().to_vec().unwrap())
                ),
            &network_info
        );
        let recipient_verification = self.recipient.as_ref().map(|rec| rec.encrypt_data(&hash).unwrap());
        let signature = match (self.sender, self.recipient) {
            (Some(sender_keys), Some(recipient_key)) => Signature::SignedPrivately {
                sender_encrypted: sender_encrypted.unwrap(),
                signature: recipient_key.encrypt_data(&sender_keys.sign(&hash)?)?,
            },
            (Some(sender_keys), None) => Signature::Signed { sender: sender_keys.get_public(), signature: sender_keys.sign(&hash)? },
            (None, _) => Signature::NotSigned,
        };
        Ok(Message {
            content: self.content,
            message_type: self.message_type.unwrap(),
            signature,
            salt,
            hash,
            recipient_verification,
            tunnel_id: self.tunnel_id,
            network_info,
        })
    }
}

#[cfg(test)]
use alloc::vec;

#[test]
fn test_hashing_message_type() {
    let msg_type_1 = MessageType::Broadcast;
    let msg_type_2 = MessageType::Service(ServiceMessageType::TunnelBuilding(TunnelPublic::new_for_test()));
    assert_eq!(msg_type_1.hash(), msg_type_1.hash());
    assert_eq!(msg_type_2.hash(), msg_type_2.hash());
    assert_ne!(msg_type_1.hash(), msg_type_2.hash())
}

#[test]
fn test_hash_message_content() {
    let content_1 = MessageContent::Plain(vec![1, 2, 4, 5]);
    let content_2 = MessageContent::Encrypted(vec![1, 2, 4, 5]);
    let content_3 = MessageContent::Plain(vec![1, 3, 4, 5]);
    assert_eq!(content_1.hash(), content_1.hash());
    assert_ne!(content_1.hash(), MessageContent::None.hash());
    assert_ne!(content_1.hash(), content_2.hash());
    assert_ne!(content_1.hash(), content_3.hash());
    assert_ne!(content_3.hash(), content_2.hash());
}