cusku/src/bot.rs

// Copyright (C) 2023-2099 The crate authors.
//
// This program is free software: you can redistribute it and/or modify it
// under the terms of the GNU Affero General Public License as published by the
// Free Software Foundation, either version 3 of the License, or (at your
// option) any later version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License
// for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

use crate::Error;
use crate::config::Config;
use crate::hooks::{Hook, IssueAction, MergeRequestAction, format_hook};

use chrono::{TimeDelta, Utc};
use log::debug;
use tokio::{signal::ctrl_c, sync::mpsc};
use xmpp::jid::{BareJid, Jid, ResourcePart};
use xmpp::parsers::message::MessageType;
use xmpp::{
	Agent, ClientBuilder, ClientFeature, ClientType, Config as AgentConfig, Event, RoomNick,
	message::send::{MessageSettings, RawMessageSettings},
	muc::room::JoinRoomSettings,
};

pub struct XmppClient {
	is_online: bool,
	agent: Agent,
	nickname: ResourcePart,
	config: Config,
	/// Keep around messages we've sent recently so we're able to prevent spamming the same type of
	/// messages
	recent_hooks: Vec<Hook>,
}

impl XmppClient {
	pub fn new(jid: BareJid, password: &str, nickname: ResourcePart, config: Config) -> XmppClient {
		let agent_config = AgentConfig {
			bookmarks_autojoin: false,
			..AgentConfig::default()
		};
		let agent = ClientBuilder::new(jid, password)
			.set_config(agent_config)
			.set_client(ClientType::Bot, "xmpp-rs")
			.set_website("https://gitlab.com/xmpp-rs/xmpp-rs")
			.set_default_nick(&nickname)
			.enable_feature(ClientFeature::JoinRooms)
			.build();

		XmppClient {
			is_online: false,
			agent,
			nickname,
			config,
			recent_hooks: Vec::new(),
		}
	}

	pub async fn next(&mut self) {
		for event in self.agent.wait_for_events().await {
			match event {
				Event::Online => {
					self.is_online = true;
					debug!("XMPP Online");
					for room in &self.config.rooms {
						self.agent
							.join_room(JoinRoomSettings {
								nick: Some(RoomNick::from_resource_ref(self.nickname.as_ref())),
								..JoinRoomSettings::new(room.clone())
							})
							.await
					}
				}
				Event::ChatMessage(_id, bare, message, _timeinfo) => {
					if !self.config.admins.contains(&bare) {
						debug!("Received chat message from {}, not in admins", bare);
						continue;
					}

					debug!("Received chat message from {}: {}", bare, message);

					match message.as_str() {
						"rejoin" => {
							for room in &self.config.rooms {
								self.agent
									.join_room(JoinRoomSettings {
										nick: Some(RoomNick::from_resource_ref(
											self.nickname.as_ref(),
										)),
										force_resync: true,
										..JoinRoomSettings::new(room.clone())
									})
									.await
							}
						}
						"version" => {
							self.agent
								.send_message(MessageSettings {
									recipient: bare,
									message: self.config.version.unwrap_or("No version specified."),
									lang: Some("en"),
								})
								.await
						}
						_ => {
							let message = "Help:\nrejoin: Re-join all rooms.\nversion: Display software version.";
							self.agent
								.send_message(MessageSettings {
									recipient: bare,
									message,
									lang: Some("en"),
								})
								.await
						}
					}
				}
				Event::Disconnected(e) => {
					self.is_online = false;
					debug!("XMPP Disconnected: {e}");
				}
				_ => {
					debug!("XMPP Event not supported")
				}
			}
		}
	}

	pub async fn receive(&mut self, mut rx: mpsc::UnboundedReceiver<Hook>) {
		loop {
			tokio::select! {
				_ = ctrl_c() => {
						return;  // Disconnecting
				},
				_ = self.next() => (),
				wh = rx.recv() => {
					if let Some(hook) = wh {
						debug!("XMPP Bot Received Hook");
						self.hook(hook).await
					}
				}
			}
		}
	}

	/// Compare incoming hook with recently sent hooks and update the list.
	/// Compare only hook type and author. If they match, check how long ago it was.
	/// Returns true if the hook is recent enough and thus not to be sent.
	fn update_recent_hooks(&mut self, new_hook: Hook) -> bool {
		// Do we need to remove/update the hook?
		let mut removal: Option<usize> = None;

		let datetime = Utc::now() - TimeDelta::minutes(5);

		// Remove expired hooks (more than time delta)
		self.recent_hooks.retain(|hook| match hook {
			Hook::MergeRequest(hook) if hook.updated_at > datetime => false,
			Hook::Issue(hook) if hook.updated_at > datetime => false,
			_ => true,
		});

		for (i, old_hook) in self.recent_hooks.iter().enumerate() {
			match (old_hook, &new_hook) {
				(_, &Hook::MergeRequest(ref new))
					if new.action != Some(MergeRequestAction::Update) =>
				{
					return false;
				}
				(_, &Hook::Issue(ref new)) if new.action != Some(IssueAction::Update) => {
					return false;
				}
				(&Hook::MergeRequest(ref old), &Hook::MergeRequest(ref new)) => {
					// Action is MergeRequestAction::Update, otherwise it would have matched the other branch
					// and the method would have returned.
					if old.id == new.id {
						if old.author.name == new.author.name {
							// If everything matches and we're still within the time frame, let the old hook
							// expire, don't update it.
							return true;
						} else {
							// The old hook either doesn't match the new author (we want to announce messages
							// from different authors), let it be replaced by the new hook.
							removal = Some(i);
							break;
						}
					}
				}
				(&Hook::Issue(ref old), &Hook::Issue(ref new)) => {
					// See the MergeRequest branch for comments
					if old.id == new.id {
						if old.author.name == new.author.name {
							return true;
						} else {
							removal = Some(i);
							break;
						}
					}
				}
				_ => (),
			}
		}

		if let Some(index) = removal {
			self.recent_hooks.swap_remove(index);
		}

		// The new hook hasn't been matched and needs to be added to the list, or it has and will
		// replace the matched hook.
		match new_hook {
			Hook::Issue(_) | Hook::MergeRequest(_) => self.recent_hooks.push(new_hook),
			_ => (),
		}

		false
	}

	pub async fn hook(&mut self, wh: Hook) {
		debug!("XMPP Bot Processing Hook");
		if self.update_recent_hooks(wh.clone()) {
			debug!("Hook already sent recently");
			return;
		}
		if let Some(display) = format_hook(&wh) {
			debug!("Hook: {}", display);
			for room in &self.config.rooms {
				self.agent
					.send_raw_message(RawMessageSettings {
						recipient: Jid::from(room.clone()),
						message_type: MessageType::Groupchat,
						message: &display,
						lang: Some("en"),
						payloads: Vec::new(),
					})
					.await
			}
		}
		debug!("XMPP Bot Processed Hook");
	}

	pub async fn disconnect(self) -> Result<(), Error> {
		log::info!("Disconnecting...");
		Ok(self.agent.disconnect().await?)
	}
}