reminder-bot/src/main.rs

use std::{convert::TryFrom, fs::File, sync::Arc, time::Duration};

use matrix_sdk::{
	events::{
		reaction::{ReactionEventContent, Relation as ReactionRelation},
		room::{
			member::MemberEventContent,
			message::{MessageEvent, MessageEventContent, MessageType},
		},
		AnyMessageEventContent, AnyStrippedStateEvent, AnySyncMessageEvent, AnySyncRoomEvent,
		StrippedStateEvent,
	},
	identifiers::{RoomId, UserId},
	Client, ClientConfig, LoopCtrl, Session, SyncSettings,
};
use tokio::signal::unix::SignalKind;

const HELP_MESSAGE_PLAIN: &str = r#"Help:
!remindme help - show this help message
!remindme <time string> - send a reminder after the amount of time
!remindme <time string>: <message> - send a reminder with a message
<time string> is a string like "1h" or "1 hour"
"#;
const HELP_MESSAGE_HTML: &str = r#"<b>Help:</b>
<ul>
<li><code>!remindme &lt;time string&gt;</code>: Send a reminder at the specified time</li>
<li><code>!remindme &lt;time string&gt;: &lt;message&gt;:</code>: Send a reminder with a message</li>
</ul>
Where <code>&lt;time string&gt;</code> is a string like <code>1 day 23 seconds</code> or <code>1h 5m</code>. The HTML part of messages is ignored.
"#;

#[tokio::main]
async fn main() {
	if let Some("login") = std::env::args().nth(1).as_deref() {
		println!("Performing interactive login");
		let path = std::env::var("SESSION_FILE").unwrap_or_else(|_| String::from("session.json"));
		let file = std::fs::File::create(&path).expect("Can't open session file");
		let session = match login().await {
			Ok(session) => session,
			Err(e) => {
				eprintln!("Login failed: {}", e);
				return;
			}
		};
		serde_json::to_writer_pretty(file, &session).unwrap();
		println!("Interactive login completed");
		return;
	}
	let client = match restore_login().await {
		Ok(client) => client,
		Err(e) => {
			eprintln!("Login failed: {}", e);
			return;
		}
	};
	sync(client).await
}

/// Constructs a client from an mxid with the appropriate configuration. A
/// storage directory will be set if the `STORAGE_DIR` environment variable is
/// set.
async fn client(user_id: UserId) -> matrix_sdk::Result<Client> {
	match std::env::var_os("STORAGE_DIR") {
		Some(dir) => {
			Client::new_from_user_id_with_config(user_id, ClientConfig::new().store_path(&dir))
				.await
		}
		None => Client::new_from_user_id(user_id).await,
	}
}

/// Log in interactively via the TTY.
async fn login() -> Result<Session, Box<dyn std::error::Error>> {
	println!("Enter username (e.g. @alice:matrix.org):");
	let mut user = String::new();
	std::io::stdin().read_line(&mut user)?;
	let user = UserId::try_from(user.trim())?;
	println!("Enter password:");
	let password = rpassword::read_password_from_tty(None)?;
	let client = client(user.clone()).await?;
	let response = client.login(user.as_str(), password.trim(), None, Some("Reminder bot")).await?;
	let session = Session {
		access_token: response.access_token,
		user_id: response.user_id,
		device_id: response.device_id,
	};
	Ok(session)
}

/// Restores a login from a session file. It gets the file path from the
/// `SESSION_PATH` environment variable if set, and defaults to looking for
/// `session.json` in the current directory otherwise.
async fn restore_login() -> Result<Client, Box<dyn std::error::Error>> {
	let path = std::env::var("SESSION_FILE").unwrap_or_else(|_| String::from("session.json"));
	let file = File::open(&path)?;
	let session: Session = serde_json::from_reader(file)?;
	let client = client(session.user_id.clone()).await?;
	client.restore_login(session).await?;
	Ok(client)
}

async fn sync(client: Client) {
	let stop = Arc::new(tokio::sync::RwLock::new(false));
	// install signal handler
	{
		let stop = stop.clone();
		tokio::task::spawn(async move {
			let mut stream = match tokio::signal::unix::signal(SignalKind::terminate()) {
				Ok(stream) => stream,
				Err(e) => {
					eprintln!("Attaching signal handler failed: {}", e);
					return;
				}
			};
			loop {
				stream.recv().await;
				*stop.write().await = true;
			}
		});
	}
	client
		.sync_with_callback(SyncSettings::new(), |response| {
			let client = client.clone();
			let stop = stop.clone();
			async move {
				for (room_id, room) in response.rooms.invite {
					for event in room.invite_state.events {
						let event = match event.deserialize() {
							Ok(event) => event,
							Err(_) => continue,
						};
						if let AnyStrippedStateEvent::RoomMember(member) = event {
							handle_invite(client.clone(), &room_id, member).await
						};
					}
				}
				for (room_id, room) in response.rooms.join {
					for event in room.timeline.events {
						// skip if serialization failed
						let event = match event.event.deserialize() {
							Ok(event) => event,
							Err(_) => continue,
						};
						// Skip if not a text message, otherwise extract the text event
						if let AnySyncRoomEvent::Message(AnySyncMessageEvent::RoomMessage(
							message,
						)) = event
						{
							handle_message(client.clone(), message.into_full_event(room_id.clone()))
								.await
						};
					}
				}
				match *stop.read().await {
					false => LoopCtrl::Continue,
					true => {
						eprintln!("Shutting down gracefully");
						LoopCtrl::Break
					}
				}
			}
		})
		.await;
}

async fn handle_invite(
	client: Client,
	room_id: &RoomId,
	member: StrippedStateEvent<MemberEventContent>,
) {
	// Check that it's an invite, and that it's for our own mxid
	if member.state_key != client.user_id().await.unwrap() {
		return;
	}
	let invited = match client.get_invited_room(room_id) {
		Some(invited) => invited,
		None => return,
	};
	// Retry the accept until an house passes because synapse is silly and sends the
	// invite before it's ready for receiving a join.
	let mut delay = std::time::Duration::from_secs(2);
	while invited.accept_invitation().await.is_err() {
		tokio::time::sleep(delay).await;
		delay *= 2;
		if delay.as_secs() > 3600 {
			return;
		}
	}
}

/// Handles a message event and sets a reminder if a correctly invoked command is found.
async fn handle_message(client: Client, message: MessageEvent) {
	let text = match message.content.msgtype {
		MessageType::Text(ref text) => text,
		_ => return,
	};
	// Strip the command prefix, skip if it isn't there
	let time = match text.body.strip_prefix("!remindme ") {
		Some(time) => time,
		None => return,
	};
	// Get the reminder message if it exists.
	let (time, reminder) = match time.split_once(": ") {
		Some((time, reminder)) => (time, Some(reminder.to_owned())),
		None => (time, None),
	};

	// Get the room so we can send messages.
	let room = client.get_joined_room(&message.room_id).unwrap();

	// Check for a help message
	if time.trim().to_ascii_lowercase().as_str() == "help" {
		let response = AnyMessageEventContent::RoomMessage(MessageEventContent::notice_reply_html(
			HELP_MESSAGE_PLAIN,
			HELP_MESSAGE_HTML,
			&message,
		));
		let _ = room.send(response, None).await;
		return;
	}

	// Parse the duration, send message if that fails.
	let duration = match humantime::parse_duration(time) {
		Ok(time) => time,
		Err(_) if time.trim() == "tomorrow" => Duration::new(60 * 60 * 24, 0),
		Err(e) => {
			let response =
				AnyMessageEventContent::RoomMessage(MessageEventContent::notice_reply_plain(
					format!("Sorry! I don't understand when you want me to remind you: {}. Say !remindme help for help.", e),
					&message,
				));
			let _ = room.send(response, None).await;
			return;
		}
	};
	// Send a reaction so the sender knows their message was received
	let remind_time = time::OffsetDateTime::now_utc() + duration;
	let _ = room
		.send(
			AnyMessageEventContent::Reaction(ReactionEventContent::new(ReactionRelation::new(
				message.event_id.clone(),
				remind_time.format("%d %b %T UTC"),
			))),
			None,
		)
		.await;
	tokio::task::spawn(async move {
		tokio::time::sleep(duration).await;
		// contruct the reminder message
		let content = match reminder {
			Some(reminder) => format!("Here's your reminder: {}", reminder),
			None => String::from("Here's your reminder!"),
		};
		let response = AnyMessageEventContent::RoomMessage(
			MessageEventContent::notice_reply_plain(content, &message),
		);
		let _ = room.send(response, None).await;
	});
}

#[test]
fn parsing_test() {
	humantime::parse_duration("1 day 23 seconds").unwrap();
}