matrix-bot/src/main.rs

#![warn(clippy::nursery, clippy::pedantic)]

use std::fmt::Display;
use std::path::PathBuf;
use std::str::FromStr;
use std::time::Duration;

use anyhow::{Context, Result};
use clap::Parser;
use enum_iterator::IntoEnumIterator;
use matrix_sdk::config::SyncSettings;
use matrix_sdk::room::{Joined, MessagesOptions, Room};
use matrix_sdk::ruma::events::room::member::StrippedRoomMemberEvent;
use matrix_sdk::ruma::events::room::message::{
    MessageType, RoomMessageEventContent, SyncRoomMessageEvent, TextMessageEventContent,
};
use matrix_sdk::ruma::events::{
    AnyMessageLikeEvent, AnyRoomEvent, MessageLikeEvent, OriginalMessageLikeEvent,
    OriginalSyncMessageLikeEvent, SyncMessageLikeEvent,
};
use matrix_sdk::ruma::OwnedUserId;
use matrix_sdk::store::CryptoStore;
use matrix_sdk::{Client, Session};
use serde::{Deserialize, Serialize};
use tracing::{error, info, Level};
use tracing_subscriber::util::SubscriberInitExt;

#[derive(Parser)]
struct Args {
    #[clap(subcommand)]
    subcommand: Command,
}

#[derive(Parser)]
enum Command {
    Login {
        user_id: OwnedUserId,
        password: String,
    },
    Run {
        #[clap(default_value = "config.toml")]
        config_path: PathBuf,
        #[clap(default_value = "crypto_data")]
        crypto_store_path: PathBuf,
    },
}

#[derive(Serialize, Deserialize)]
struct TomlConfig {
    session: Session,
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_max_level(Level::INFO)
        .finish()
        .init();
    let args = Args::parse();

    match args.subcommand {
        Command::Login { user_id, password } => handle_login(user_id, &password).await,
        Command::Run {
            config_path,
            crypto_store_path,
        } => handle_run(config_path, crypto_store_path).await,
    }
}

async fn handle_login(user_id: OwnedUserId, password: &str) -> Result<()> {
    let client = Client::builder().user_id(&user_id).build().await?;
    let resp = client
        .login(user_id.localpart(), password, None, None)
        .await?;

    info!("Config generated. Please store this somewhere (like config.toml)!");
    println!(
        "\n\n{}",
        toml::to_string_pretty(&TomlConfig {
            session: Session {
                access_token: resp.access_token,
                device_id: resp.device_id,
                user_id,
            }
        })?
    );

    Ok(())
}

async fn handle_run(config_path: PathBuf, crypto_store_path: PathBuf) -> Result<()> {
    let data = std::fs::read_to_string(config_path)?;
    let config: TomlConfig = toml::from_str(&data)?;

    let client = Client::builder()
        .user_id(&config.session.user_id)
        .crypto_store(Box::new(CryptoStore::open_with_passphrase(
            crypto_store_path,
            None,
        )?))
        .build()
        .await
        .context("s")?;
    client.restore_login(config.session.clone()).await?;

    info!(user = %config.session.user_id, device_id = %config.session.device_id, "Logged in");

    let sync_settings = SyncSettings::default();

    // Sync once before registering events to not double-send events on boot.
    client.sync_once(sync_settings.clone()).await?;

    client
        .register_event_handler(auto_join)
        .await
        .register_event_handler(on_room_message)
        .await;

    // Syncing is important to synchronize the client state with the server.
    // This method will never return.
    client.sync(sync_settings).await;

    Ok(())
}

async fn auto_join(room_member: StrippedRoomMemberEvent, client: Client, room: Room) -> Result<()> {
    if Some(room_member.sender) == client.user_id().await {
        return Ok(());
    }

    if let Room::Invited(room) = room {
        info!("Autojoining room {}", room.room_id());
        let mut delay = 2;

        while let Err(err) = room.accept_invitation().await {
            // retry autojoin due to synapse sending invites, before the
            // invited user can join for more information see
            // https://github.com/matrix-org/synapse/issues/4345
            error!(
                room_id = ?room.room_id(),
                ?err,
                "Failed to join room, retrying in {delay}s",
            );

            tokio::time::sleep(Duration::from_secs(delay)).await;
            delay *= 2;

            if delay > 3600 {
                error!(
                    room_id = ?room.room_id(),
                    ?err,
                    "Failed to join room. Giving up.",
                );
                break;
            }
        }

        info!(room_id = ?room.room_id(), "Successfully joined room.");
    }

    Ok(())
}

async fn on_room_message(event: SyncRoomMessageEvent, room: Room, client: Client) -> Result<()> {
    // Ignore messages sent from self.
    if client.user_id().await.as_deref() == Some(event.sender()) {
        return Ok(());
    }

    if let Room::Joined(room) = room {
        parse_message(event, room).await?;
    }

    Ok(())
}

async fn parse_message(event: SyncRoomMessageEvent, room: Joined) -> Result<()> {
    room.read_receipt(event.event_id()).await?;

    let message = match &event {
        SyncMessageLikeEvent::Original(OriginalSyncMessageLikeEvent {
            content:
                RoomMessageEventContent {
                    msgtype: MessageType::Text(TextMessageEventContent { body, .. }),
                    ..
                },
            ..
        }) => body.clone(),
        _ => return Ok(()),
    };

    if let Some(message) = message.strip_prefix('!') {
        let (command, args) = message.split_once(' ').unwrap_or((message, ""));
        match BotCommand::from_str(command).ok() {
            Some(BotCommand::Source) => source(room).await?,
            Some(BotCommand::Uwu) => uwuify(room, args).await?,
            Some(BotCommand::Ping) => ping(room).await?,
            Some(BotCommand::Help) => help(room).await?,
            _ => unsupported_command(event, room, command).await?,
        }
    }

    Ok(())
}

#[derive(IntoEnumIterator)]
enum BotCommand {
    Help,
    Source,
    Uwu,
    Ping,
}

impl BotCommand {
    const fn help_text(&self) -> &'static str {
        match self {
            Self::Source => "Links to the source code for this bot.",
            Self::Uwu => "Uwuifies your message.",
            Self::Ping => "Pong!",
            Self::Help => "Prints this help.",
        }
    }
}

impl FromStr for BotCommand {
    type Err = ();

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "source" => Ok(Self::Source),
            "uwu" => Ok(Self::Uwu),
            "ping" => Ok(Self::Ping),
            "help" => Ok(Self::Help),
            _ => Err(()),
        }
    }
}

impl Display for BotCommand {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BotCommand::Source => write!(f, "source"),
            BotCommand::Uwu => write!(f, "uwu"),
            BotCommand::Ping => write!(f, "ping"),
            BotCommand::Help => write!(f, "help"),
        }
    }
}

async fn source(room: Joined) -> Result<()> {
    let content = RoomMessageEventContent::text_plain("https://git.eddie.sh/edward/matrix-bot");
    room.send(content, None).await?;
    Ok(())
}

async fn uwuify(room: Joined, message: &str) -> Result<()> {
    let message = if message.is_empty() {
        if let Ok(msg) = get_previous_message(&room).await {
            uwuifier::uwuify_str_sse(&msg.body)
        } else {
            "uwu".to_owned()
        }
    } else {
        uwuifier::uwuify_str_sse(message)
    };
    let content = RoomMessageEventContent::text_plain(message);
    room.send(content, None).await?;
    Ok(())
}

async fn get_previous_message(room: &Joined) -> Result<TextMessageEventContent> {
    let last_prev_batch = (**room)
        .last_prev_batch()
        .context("No previous batch key while trying to find previous message")?;
    let request = MessagesOptions::backward(&last_prev_batch);
    let messages = room.messages(request).await?;

    messages
        .chunk
        .iter()
        .filter_map(|msg| msg.event.deserialize().ok())
        .find_map(|msg| {
            if let AnyRoomEvent::MessageLike(AnyMessageLikeEvent::RoomMessage(
                MessageLikeEvent::Original(OriginalMessageLikeEvent {
                    content:
                        RoomMessageEventContent {
                            msgtype: MessageType::Text(content),
                            ..
                        },
                    ..
                }),
            )) = msg
            {
                Some(content)
            } else {
                None
            }
        })
        .context("Failed to find previous message")
}

async fn ping(room: Joined) -> Result<()> {
    let content = RoomMessageEventContent::text_plain("Pong!");
    room.send(content, None).await?;
    Ok(())
}

async fn help(room: Joined) -> Result<()> {
    let mut msg = "List of commands:\n\n".to_owned();
    for command in BotCommand::into_enum_iter() {
        msg.push_str(&command.to_string());
        msg.push_str(" - ");
        msg.push_str(command.help_text());
        msg.push('\n');
    }
    let content = RoomMessageEventContent::notice_plain(msg);
    room.send(content, None).await?;
    Ok(())
}

async fn unsupported_command(
    event: SyncRoomMessageEvent,
    room: Joined,
    command: &str,
) -> Result<()> {
    // let resp = .into_full_event(room.room_id().to_owned());
    let event = if let SyncRoomMessageEvent::Original(event) = event {
        event.into_full_event(room.room_id().to_owned())
    } else {
        return Ok(());
    };
    let content =
        RoomMessageEventContent::notice_reply_plain(format!("Unknown command `{command}`"), &event);
    room.send(content, None).await?;
    Ok(())
}