matrix-lol/src/matrix.rs

use matrix_sdk::{
    config::SyncSettings, ruma::api::client::filter::FilterDefinition, Client, Session,
};
use rand::{distributions::Alphanumeric, thread_rng, Rng};
use serde::{Deserialize, Serialize};
use std::{
    io::{self, Write},
    path::{Path, PathBuf},
};
use tokio::fs;

#[derive(Debug, Serialize, Deserialize)]
pub struct ClientSession {
    homeserver: String,
    db_path: PathBuf,
    passphrase: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct FullSession {
    client_session: ClientSession,
    user_session: Session,
    #[serde(skip_serializing_if = "Option::is_none")]
    sync_token: Option<String>,
}

pub async fn restore_session(session_file: &Path) -> anyhow::Result<(Client, Option<String>)> {
    let serialized_session = fs::read_to_string(session_file).await?;
    let FullSession {
        client_session,
        user_session,
        sync_token,
    } = serde_json::from_str(&serialized_session)?;

    let client = Client::builder()
        .homeserver_url(client_session.homeserver)
        .sled_store(client_session.db_path, Some(&client_session.passphrase))
        .unwrap()
        .build()
        .await?;

    println!("[*] Restoring session for {}…", user_session.user_id);

    client.restore_login(user_session).await?;

    Ok((client, sync_token))
}

pub async fn login(data_dir: &Path, session_file: &Path) -> anyhow::Result<Client> {
    println!("[*] No previous session found, logging in…");

    let (client, client_session) = build_client(data_dir).await?;

    loop {
        let username = rpassword::prompt_password("User> ").unwrap();
        let password = rpassword::prompt_password("Password> ").unwrap();

        match client
            .login_username(&username, &password)
            .initial_device_display_name("vcxz bot")
            .send()
            .await
        {
            Ok(_) => {
                println!("[*] Logged in as {username}");
                break;
            }
            Err(error) => {
                println!("[!] Error logging in: {error}");
                println!("[!] Please try again\n");
            }
        }
    }

    let user_session = client
        .session()
        .expect("A logged-in client should have a session");
    let serialized_session = serde_json::to_string(&FullSession {
        client_session,
        user_session,
        sync_token: None,
    })?;
    fs::write(session_file, serialized_session).await?;
    Ok(client)
}

/// Build a new client.
pub async fn build_client(data_dir: &Path) -> anyhow::Result<(Client, ClientSession)> {
    let mut rng = thread_rng();
    
    let db_subfolder: String = (&mut rng)
        .sample_iter(Alphanumeric)
        .take(7)
        .map(char::from)
        .collect();
    let db_path = data_dir.join(db_subfolder);

    let passphrase: String = (&mut rng)
        .sample_iter(Alphanumeric)
        .take(32)
        .map(char::from)
        .collect();

    loop {
        let mut homeserver = String::new();

        print!("Homeserver> ");
        io::stdout().flush().expect("Unable to write to stdout");
        io::stdin()
            .read_line(&mut homeserver)
            .expect("Unable to read user input");

        println!("\n[*] Checking homeserver…");

        match Client::builder()
            .homeserver_url(&homeserver)
            .sled_store(&db_path, Some(&passphrase))
            .unwrap()
            .build()
            .await
        {
            Ok(client) => {
                return Ok((
                    client,
                    ClientSession {
                        homeserver,
                        db_path,
                        passphrase,
                    },
                ))
            }
            Err(error) => match &error {
                matrix_sdk::ClientBuildError::AutoDiscovery(_)
                | matrix_sdk::ClientBuildError::Url(_)
                | matrix_sdk::ClientBuildError::Http(_) => {
                    println!("[!] Error checking the homeserver: {error}");
                    println!("[!] Please try again\n");
                }
                _ => {
                    return Err(error.into());
                }
            },
        }
    }
}

pub async fn sync<'a>(
    client: Client,
    initial_sync_token: Option<String>,
) -> anyhow::Result<(Client, SyncSettings<'a>)> {
    println!("[*] Initial sync...");

    let filter = FilterDefinition::empty();

    let mut sync_settings = SyncSettings::default().filter(filter.into());

    if let Some(sync_token) = initial_sync_token {
        sync_settings = sync_settings.token(sync_token);
    }

    loop {
        match client.sync_once(sync_settings.clone()).await {
            Ok(response) => {
                sync_settings = sync_settings.token(response.next_batch.clone());
                persist_sync_token(response.next_batch).await?;
                break;
            }
            Err(error) => {
                println!("[!] An error occurred during initial sync: {error}");
                println!("[!] Trying again…");
            }
        }
    }

    println!("[*] The bot is ready!");

    Ok((client, sync_settings))
}

pub async fn persist_sync_token(sync_token: String) -> anyhow::Result<()> {
    let data_dir = dirs::data_dir()
        .expect("no data_dir directory found")
        .join("vcxz_bot");
    let session_file = data_dir.join("session");

    let serialized_session = fs::read_to_string(&session_file).await?;
    let mut full_session: FullSession = serde_json::from_str(&serialized_session)?;

    full_session.sync_token = Some(sync_token);
    let serialized_session = serde_json::to_string(&full_session)?;
    fs::write(session_file, serialized_session).await?;

    Ok(())
}