luz/src/stanza/stream.rs

use std::collections::{HashMap, HashSet};

use peanuts::element::{Content, FromElement, IntoElement, NamespaceDeclaration};
use peanuts::XML_NS;
use peanuts::{element::Name, Element};

use crate::{Error, JID};

use super::starttls::StartTls;

pub const XMLNS: &str = "http://etherx.jabber.org/streams";
pub const XMLNS_CLIENT: &str = "jabber:client";

// MUST be qualified by stream namespace
// #[derive(XmlSerialize, XmlDeserialize)]
// #[peanuts(xmlns = XMLNS)]
#[derive(Debug)]
pub struct Stream {
    pub from: Option<JID>,
    to: Option<JID>,
    id: Option<String>,
    version: Option<String>,
    // TODO: lang enum
    lang: Option<String>,
    // #[peanuts(content)]
    // content: Message,
}

impl FromElement for Stream {
    fn from_element(element: Element) -> peanuts::Result<Self> {
        let Name {
            namespace,
            local_name,
        } = element.name;
        if namespace.as_deref() == Some(XMLNS) && &local_name == "stream" {
            let (mut from, mut to, mut id, mut version, mut lang) = (None, None, None, None, None);
            for (name, value) in element.attributes {
                match (name.namespace.as_deref(), name.local_name.as_str()) {
                    (None, "from") => from = Some(value.try_into()?),
                    (None, "to") => to = Some(value.try_into()?),
                    (None, "id") => id = Some(value),
                    (None, "version") => version = Some(value),
                    (Some(XML_NS), "lang") => lang = Some(value),
                    _ => return Err(peanuts::Error::UnexpectedAttribute(name)),
                }
            }
            return Ok(Stream {
                from,
                to,
                id,
                version,
                lang,
            });
        } else {
            return Err(peanuts::Error::IncorrectName(Name {
                namespace,
                local_name,
            }));
        }
    }
}

impl IntoElement for Stream {
    fn into_element(&self) -> Element {
        let mut namespace_declaration_overrides = HashSet::new();
        namespace_declaration_overrides.insert(NamespaceDeclaration {
            prefix: Some("stream".to_string()),
            namespace: XMLNS.to_string(),
        });
        namespace_declaration_overrides.insert(NamespaceDeclaration {
            prefix: None,
            // TODO: don't default to client
            namespace: XMLNS_CLIENT.to_string(),
        });

        let mut attributes = HashMap::new();
        self.from.as_ref().map(|from| {
            attributes.insert(
                Name {
                    namespace: None,
                    local_name: "from".to_string(),
                },
                from.to_string(),
            );
        });
        self.to.as_ref().map(|to| {
            attributes.insert(
                Name {
                    namespace: None,
                    local_name: "to".to_string(),
                },
                to.to_string(),
            );
        });
        self.id.as_ref().map(|id| {
            attributes.insert(
                Name {
                    namespace: None,
                    local_name: "id".to_string(),
                },
                id.clone(),
            );
        });
        self.version.as_ref().map(|version| {
            attributes.insert(
                Name {
                    namespace: None,
                    local_name: "version".to_string(),
                },
                version.clone(),
            );
        });
        self.lang.as_ref().map(|lang| {
            attributes.insert(
                Name {
                    namespace: Some(XML_NS.to_string()),
                    local_name: "lang".to_string(),
                },
                lang.to_string(),
            );
        });

        Element {
            name: Name {
                namespace: Some(XMLNS.to_string()),
                local_name: "stream".to_string(),
            },
            namespace_declaration_overrides,
            attributes,
            content: Vec::new(),
        }
    }
}

impl<'s> Stream {
    pub fn new(
        from: Option<JID>,
        to: Option<JID>,
        id: Option<String>,
        version: Option<String>,
        lang: Option<String>,
    ) -> Self {
        Self {
            from,
            to,
            id,
            version,
            lang,
        }
    }

    /// For initial stream headers, the initiating entity SHOULD include the 'xml:lang' attribute.
    /// For privacy, it is better to not set `from` when sending a client stanza over an unencrypted connection.
    pub fn new_client(from: Option<JID>, to: JID, id: Option<String>, lang: String) -> Self {
        Self {
            from,
            to: Some(to),
            id,
            version: Some("1.0".to_string()),
            lang: Some(lang),
        }
    }
}

#[derive(Debug)]
pub struct Features {
    features: Vec<Feature>,
}

impl IntoElement for Features {
    fn into_element(&self) -> Element {
        let mut content = Vec::new();
        for feature in &self.features {
            match feature {
                Feature::StartTls(start_tls) => {
                    content.push(Content::Element(start_tls.into_element()))
                }
                Feature::Sasl => {}
                Feature::Bind => {}
                Feature::Unknown => {}
            }
        }
        Element {
            name: Name {
                namespace: Some(XMLNS.to_string()),
                local_name: "features".to_string(),
            },
            namespace_declaration_overrides: HashSet::new(),
            attributes: HashMap::new(),
            content,
        }
    }
}

impl FromElement for Features {
    fn from_element(element: Element) -> peanuts::Result<Self> {
        let Name {
            namespace,
            local_name,
        } = element.name;
        if namespace.as_deref() == Some(XMLNS) && &local_name == "features" {
            let mut features = Vec::new();
            for feature in element.content {
                match feature {
                    Content::Element(element) => {
                        if let Ok(start_tls) = FromElement::from_element(element) {
                            features.push(Feature::StartTls(start_tls))
                        } else {
                            features.push(Feature::Unknown)
                        }
                    }
                    c => return Err(peanuts::Error::UnexpectedContent(c.clone())),
                }
            }
            return Ok(Self { features });
        } else {
            return Err(peanuts::Error::IncorrectName(Name {
                namespace,
                local_name,
            }));
        }
    }
}

#[derive(Debug)]
pub enum Feature {
    StartTls(StartTls),
    Sasl,
    Bind,
    Unknown,
}