use futures::stream::StreamExt; use std::env::args; use std::fs::{create_dir_all, File}; use std::io::{self, Write}; use std::process::exit; use std::str::FromStr; use tokio_xmpp::{Client, Stanza}; use xmpp_parsers::{ avatar::{Data as AvatarData, Metadata as AvatarMetadata}, caps::{compute_disco, hash_caps, Caps}, disco::{DiscoInfoQuery, DiscoInfoResult, Feature, Identity}, hashes::Algo, iq::{Iq, IqType}, jid::{BareJid, Jid}, ns, presence::{Presence, Type as PresenceType}, pubsub::{ event::PubSubEvent, pubsub::{Items, PubSub}, NodeName, }, stanza_error::{DefinedCondition, ErrorType, StanzaError}, }; #[tokio::main] async fn main() { env_logger::init(); let args: Vec = args().collect(); if args.len() != 3 { println!("Usage: {} ", args[0]); exit(1); } let jid = BareJid::from_str(&args[1]).expect(&format!("Invalid JID: {}", &args[1])); let password = args[2].clone(); // Client instance let mut client = Client::new(jid.clone(), password); let disco_info = make_disco(); // Main loop, processes events let mut wait_for_stream_end = false; let mut stream_ended = false; while !stream_ended { if let Some(event) = client.next().await { if wait_for_stream_end { /* Do nothing */ } else if event.is_online() { println!("Online!"); let caps = get_disco_caps(&disco_info, "https://gitlab.com/xmpp-rs/tokio-xmpp"); let presence = make_presence(caps); client.send_stanza(presence.into()).await.unwrap(); } else if let Some(stanza) = event.into_stanza() { match stanza { Stanza::Iq(iq) => { if let IqType::Get(payload) = iq.payload { if payload.is("query", ns::DISCO_INFO) { let query = DiscoInfoQuery::try_from(payload); match query { Ok(query) => { let mut disco = disco_info.clone(); disco.node = query.node; let iq = Iq::from_result(iq.id, Some(disco)) .with_to(iq.from.unwrap()); client.send_stanza(iq.into()).await.unwrap(); } Err(err) => client .send_stanza( make_error( iq.from.unwrap(), iq.id, ErrorType::Modify, DefinedCondition::BadRequest, &format!("{}", err), ) .into(), ) .await .unwrap(), } } else { // We MUST answer unhandled get iqs with a service-unavailable error. client .send_stanza( make_error( iq.from.unwrap(), iq.id, ErrorType::Cancel, DefinedCondition::ServiceUnavailable, "No handler defined for this kind of iq.", ) .into(), ) .await .unwrap(); } } else if let IqType::Result(Some(payload)) = iq.payload { if payload.is("pubsub", ns::PUBSUB) { let pubsub = PubSub::try_from(payload).unwrap(); let from = iq.from.clone().unwrap_or(jid.clone().into()); handle_iq_result(pubsub, &from); } } else if let IqType::Set(_) = iq.payload { // We MUST answer unhandled set iqs with a service-unavailable error. client .send_stanza( make_error( iq.from.unwrap(), iq.id, ErrorType::Cancel, DefinedCondition::ServiceUnavailable, "No handler defined for this kind of iq.", ) .into(), ) .await .unwrap(); } } Stanza::Message(message) => { let from = message.from.clone().unwrap(); if let Some(body) = message.get_best_body(vec!["en"]) { if body.0 == "die" { println!("Secret die command triggered by {}", from); wait_for_stream_end = true; client.send_end().await.unwrap(); } } for child in message.payloads { if child.is("event", ns::PUBSUB_EVENT) { let event = PubSubEvent::try_from(child).unwrap(); if let PubSubEvent::PublishedItems { node, items } = event { if node.0 == ns::AVATAR_METADATA { for item in items.into_iter() { let payload = item.payload.clone().unwrap(); if payload.is("metadata", ns::AVATAR_METADATA) { // TODO: do something with these metadata. let _metadata = AvatarMetadata::try_from(payload).unwrap(); println!( "{} has published an avatar, downloading...", from.clone() ); let iq = download_avatar(from.clone()); client.send_stanza(iq.into()).await.unwrap(); } } } } } } } // Nothing to do here. Stanza::Presence(_) => (), } } } else { println!("stream_ended"); stream_ended = true; } } } fn make_error( to: Jid, id: String, type_: ErrorType, condition: DefinedCondition, text: &str, ) -> Iq { let error = StanzaError::new(type_, condition, "en", text); Iq::from_error(id, error).with_to(to) } fn make_disco() -> DiscoInfoResult { let identities = vec![Identity::new("client", "bot", "en", "tokio-xmpp")]; let features = vec![ Feature::new(ns::DISCO_INFO), Feature::new(format!("{}+notify", ns::AVATAR_METADATA)), ]; DiscoInfoResult { node: None, identities, features, extensions: vec![], } } fn get_disco_caps(disco: &DiscoInfoResult, node: &str) -> Caps { let caps_data = compute_disco(disco); let hash = hash_caps(&caps_data, Algo::Sha_1).unwrap(); Caps::new(node, hash) } // Construct a fn make_presence(caps: Caps) -> Presence { let mut presence = Presence::new(PresenceType::None).with_priority(-1); presence.set_status("en", "Downloading avatars."); presence.add_payload(caps); presence } fn download_avatar(from: Jid) -> Iq { Iq::from_get( "coucou", PubSub::Items(Items { max_items: None, node: NodeName(String::from(ns::AVATAR_DATA)), subid: None, items: Vec::new(), }), ) .with_to(from) } fn handle_iq_result(pubsub: PubSub, from: &Jid) { if let PubSub::Items(items) = pubsub { if items.node.0 == ns::AVATAR_DATA { for item in items.items { match (item.id.clone(), item.payload.clone()) { (Some(id), Some(payload)) => { let data = AvatarData::try_from(payload).unwrap(); save_avatar(from, id.0, &data.data).unwrap(); } _ => {} } } } } } // TODO: may use tokio? fn save_avatar(from: &Jid, id: String, data: &[u8]) -> io::Result<()> { let directory = format!("data/{}", from); let filename = format!("data/{}/{}", from, id); println!( "Saving avatar from {} to {}.", from, filename ); create_dir_all(directory)?; let mut file = File::create(filename)?; file.write_all(data) }