1use futures::stream::StreamExt;
2use std::env::args;
3use std::fs::{create_dir_all, File};
4use std::io::{self, Write};
5use std::process::exit;
6use std::str::FromStr;
7use tokio_xmpp::rustls;
8use tokio_xmpp::{Client, Stanza};
9use xmpp_parsers::{
10 avatar::{Data as AvatarData, Metadata as AvatarMetadata},
11 caps::{compute_disco, hash_caps, Caps},
12 disco::{DiscoInfoQuery, DiscoInfoResult, Feature, Identity},
13 hashes::Algo,
14 iq::Iq,
15 jid::{BareJid, Jid},
16 ns,
17 presence::{Presence, Type as PresenceType},
18 pubsub::{
19 self,
20 pubsub::{Items, PubSub},
21 NodeName,
22 },
23 stanza_error::{DefinedCondition, ErrorType, StanzaError},
24};
25
26#[tokio::main]
27async fn main() {
28 env_logger::init();
29
30 rustls::crypto::aws_lc_rs::default_provider()
31 .install_default()
32 .expect("failed to install rustls crypto provider");
33
34 let args: Vec<String> = args().collect();
35 if args.len() != 3 {
36 println!("Usage: {} <jid> <password>", args[0]);
37 exit(1);
38 }
39 let jid = BareJid::from_str(&args[1]).expect(&format!("Invalid JID: {}", &args[1]));
40 let password = args[2].clone();
41
42 // Client instance
43 let mut client = Client::new(jid.clone(), password);
44
45 let disco_info = make_disco();
46
47 // Main loop, processes events
48 while let Some(event) = client.next().await {
49 if event.is_online() {
50 println!("Online!");
51
52 let caps = get_disco_caps(&disco_info, "https://gitlab.com/xmpp-rs/tokio-xmpp");
53 let presence = make_presence(caps);
54 client.send_stanza(presence.into()).await.unwrap();
55 } else if let Some(stanza) = event.into_stanza() {
56 match stanza {
57 Stanza::Iq(Iq::Get {
58 payload, id, from, ..
59 }) => {
60 if payload.is("query", ns::DISCO_INFO) {
61 let query = DiscoInfoQuery::try_from(payload);
62 match query {
63 Ok(query) => {
64 let mut disco = disco_info.clone();
65 disco.node = query.node;
66 let iq = Iq::from_result(id, Some(disco)).with_to(from.unwrap());
67 client.send_stanza(iq.into()).await.unwrap();
68 }
69 Err(err) => {
70 client
71 .send_stanza(
72 make_error(
73 from.unwrap(),
74 id,
75 ErrorType::Modify,
76 DefinedCondition::BadRequest,
77 &format!("{}", err),
78 )
79 .into(),
80 )
81 .await
82 .unwrap();
83 }
84 }
85 } else {
86 // We MUST answer unhandled get iqs with a service-unavailable error.
87 client
88 .send_stanza(
89 make_error(
90 from.unwrap(),
91 id,
92 ErrorType::Cancel,
93 DefinedCondition::ServiceUnavailable,
94 "No handler defined for this kind of iq.",
95 )
96 .into(),
97 )
98 .await
99 .unwrap();
100 }
101 }
102 Stanza::Iq(Iq::Result {
103 payload: Some(payload),
104 from,
105 ..
106 }) => {
107 if payload.is("pubsub", ns::PUBSUB) {
108 let pubsub = PubSub::try_from(payload).unwrap();
109 let from = from.unwrap_or(jid.clone().into());
110 handle_iq_result(pubsub, &from);
111 }
112 }
113 Stanza::Iq(Iq::Set { from, id, .. }) => {
114 // We MUST answer unhandled set iqs with a service-unavailable error.
115 client
116 .send_stanza(
117 make_error(
118 from.unwrap(),
119 id,
120 ErrorType::Cancel,
121 DefinedCondition::ServiceUnavailable,
122 "No handler defined for this kind of iq.",
123 )
124 .into(),
125 )
126 .await
127 .unwrap();
128 }
129 Stanza::Iq(Iq::Error { .. }) | Stanza::Iq(Iq::Result { payload: None, .. }) => (),
130 Stanza::Message(message) => {
131 let from = message.from.clone().unwrap();
132 if let Some(body) = message.get_best_body(vec!["en"]) {
133 if body.0 == "die" {
134 println!("Secret die command triggered by {}", from);
135 break;
136 }
137 }
138 for child in message.payloads {
139 if child.is("event", ns::PUBSUB_EVENT) {
140 let event = pubsub::Event::try_from(child).unwrap();
141 if let pubsub::event::Payload::Items {
142 node,
143 published,
144 retracted: _,
145 } = event.payload
146 {
147 if node.0 == ns::AVATAR_METADATA {
148 for item in published.into_iter() {
149 let payload = item.payload.clone().unwrap();
150 if payload.is("metadata", ns::AVATAR_METADATA) {
151 // TODO: do something with these metadata.
152 let _metadata =
153 AvatarMetadata::try_from(payload).unwrap();
154 println!(
155 "[1m{}[0m has published an avatar, downloading...",
156 from.clone()
157 );
158 let iq = download_avatar(from.clone());
159 client.send_stanza(iq.into()).await.unwrap();
160 }
161 }
162 }
163 }
164 }
165 }
166 }
167 // Nothing to do here.
168 Stanza::Presence(_) => (),
169 }
170 }
171 }
172}
173
174fn make_error(
175 to: Jid,
176 id: String,
177 type_: ErrorType,
178 condition: DefinedCondition,
179 text: &str,
180) -> Iq {
181 let error = StanzaError::new(type_, condition, "en", text);
182 Iq::from_error(id, error).with_to(to)
183}
184
185fn make_disco() -> DiscoInfoResult {
186 let identities = vec![Identity::new("client", "bot", "en", "tokio-xmpp")];
187 let features = vec![
188 Feature::new(ns::DISCO_INFO),
189 Feature::new(format!("{}+notify", ns::AVATAR_METADATA)),
190 ];
191 DiscoInfoResult {
192 node: None,
193 identities,
194 features,
195 extensions: vec![],
196 }
197}
198
199fn get_disco_caps(disco: &DiscoInfoResult, node: &str) -> Caps {
200 let caps_data = compute_disco(disco);
201 let hash = hash_caps(&caps_data, Algo::Sha_1).unwrap();
202 Caps::new(node, hash)
203}
204
205// Construct a <presence/>
206fn make_presence(caps: Caps) -> Presence {
207 let mut presence = Presence::new(PresenceType::None).with_priority(-1);
208 presence.set_status("en", "Downloading avatars.");
209 presence.add_payload(caps);
210 presence
211}
212
213fn download_avatar(from: Jid) -> Iq {
214 Iq::from_get(
215 "coucou",
216 PubSub::Items(Items {
217 max_items: None,
218 node: NodeName(String::from(ns::AVATAR_DATA)),
219 subid: None,
220 items: Vec::new(),
221 }),
222 )
223 .with_to(from)
224}
225
226fn handle_iq_result(pubsub: PubSub, from: &Jid) {
227 if let PubSub::Items(items) = pubsub {
228 if items.node.0 == ns::AVATAR_DATA {
229 for item in items.items {
230 match (item.id.clone(), item.payload.clone()) {
231 (Some(id), Some(payload)) => {
232 let data = AvatarData::try_from(payload).unwrap();
233 save_avatar(from, id.0, &data.data).unwrap();
234 }
235 _ => {}
236 }
237 }
238 }
239 }
240}
241
242// TODO: may use tokio?
243fn save_avatar(from: &Jid, id: String, data: &[u8]) -> io::Result<()> {
244 let directory = format!("data/{}", from);
245 let filename = format!("data/{}/{}", from, id);
246 println!(
247 "Saving avatar from [1m{}[0m to [4m{}[0m.",
248 from, filename
249 );
250 create_dir_all(directory)?;
251 let mut file = File::create(filename)?;
252 file.write_all(data)
253}