matrix-media-event-decrypt/src/main.rs

use std::fs;
use std::process::exit;
use std::io::Read;

use anyhow::Result;

use clap::Parser;
use log::{info, debug, error};
use serde::Deserialize;
use rustyline_async::{Readline, ReadlineEvent};

use matrix_sdk::{matrix_auth::{MatrixSession, MatrixSessionTokens}, ruma::{api::client::{authenticated_media, media}, device_id, user_id}, Client, SessionMeta};
use matrix_sdk::ruma::events::room::{MediaSource, message::{MessageType, RoomMessageEventContent}};

// NOTE: This has been stolen from the matrix-rust-sdk to patch out checks for authenticated media
async fn get_media_content(
    client: &Client,
    source: &MediaSource,
    use_auth: bool
) -> matrix_sdk::Result<Vec<u8>> {
    let content: Vec<u8> = match &source {
        MediaSource::Encrypted(file) => {
            let content = if use_auth {
                let request = authenticated_media::get_content::v1::Request::from_uri(&file.url)?;
                client.send(request, None).await?.file
            } else {
                #[allow(deprecated)]
                let request = media::get_content::v3::Request::from_url(&file.url)?;
                client.send(request, None).await?.file
            };

            let content = {
                debug!("Decrypting file");

                let content_len = content.len();
                let mut cursor = std::io::Cursor::new(content);
                let mut reader = match matrix_sdk::crypto::AttachmentDecryptor::new(
                    &mut cursor,
                    file.as_ref().clone().into(),
                ) {
                    Ok(reader) => reader,
                    Err(error) => {
                        error!("Unable to decrypt: {}", error);
                        exit(1);
                    }
                };

                // Encrypted size should be the same as the decrypted size,
                // rounded up to a cipher block.
                let mut decrypted = Vec::with_capacity(content_len);

                reader.read_to_end(&mut decrypted)?;

                decrypted
            };

            content
        }
        MediaSource::Plain(uri) => {
            debug!("File doesn't appear to be encrypted");

            if use_auth {
                let request = authenticated_media::get_content::v1::Request::from_uri(uri)?;
                client.send(request, None).await?.file
            } else {
                #[allow(deprecated)]
                let request = media::get_content::v3::Request::from_url(uri)?;
                client.send(request, None).await?.file
            }
        }
    };

    Ok(content)
}

async fn process_media(source: MediaSource, body: Option<String>, server: Option<String>) -> Result<()> {
    let uri = match source.clone() {
        MediaSource::Plain(uri) => uri,
        MediaSource::Encrypted(encrypted) => encrypted.url,
    };

    let server_name = uri.server_name()?;
    let media_id = uri.media_id()?;

    let client = match server {
        Some(ref url) => Client::builder().server_name_or_homeserver_url(url),
        None => Client::builder().server_name(server_name),
    }.build().await?;

    debug!("Downloading {} from {}", uri, client.homeserver());

    let content = match get_media_content(&client, &source, false).await {
        Ok(content) => content,
        Err(error) => {
            let error_is_needs_authentication = match &error {
                matrix_sdk::Error::Http(http) => match http {
                    matrix_sdk::HttpError::IntoHttp(into_http) => {
                        match into_http {
                            matrix_sdk::ruma::api::error::IntoHttpError::NeedsAuthentication => true,
                            _ => false,
                        }
                    },
                    _ => false
                },
                _ => false,
            };

            if error_is_needs_authentication {
                info!("The server did not allow downloading without authentication, please enter an access token:");

                let token = rpassword::prompt_password("> ")?;

                client.restore_session(MatrixSession {
                    meta: SessionMeta {
                        user_id: user_id!("@unused:example.com").to_owned(),
                        device_id: device_id!("UNUSED").to_owned(),
                    },
                    tokens: MatrixSessionTokens {
                        access_token: token,
                        refresh_token: None,
                    },
                }).await?;

                get_media_content(&client, &source, false).await
                    .map_err(|error| anyhow::Error::msg(error.to_string()))?
            } else {
                return Err(anyhow::Error::msg(error.to_string()));
            }
        }
    };

    let file_name = match body {
        Some(body) => format!("{server_name}-{media_id}-{}", body),
        None => format!("{server_name}-{media_id}")
    };

    info!("Saving to {file_name}");

    fs::write(file_name, content)?;

    Ok(())
}

#[derive(Deserialize, Debug)]
struct DecryptedEventSource {
    content: RoomMessageEventContent,
}

#[derive(clap::Parser)]
#[command(author, version, about, long_about = None)]
struct Arguments {
    /// Homeserver URL or Name
    #[clap(long, env)]
    server: Option<String>,
}

#[tokio::main]
async fn main() -> Result<()> {
    env_logger::init_from_env(env_logger::Env::default().filter_or(env_logger::DEFAULT_FILTER_ENV, "matrix_media_event_decrypt=debug"));

    let Arguments {
        server
    } = Arguments::parse();

    info!("Enter decrypted event source, confirm with Ctrl+D");

    let (mut readline, mut _stdout) = Readline::new("> ".to_owned()).unwrap();
    let mut lines = Vec::new();

    loop {
        match readline.readline().await {
            Ok(ReadlineEvent::Line(line)) => {
                let line = line.trim();
                if line.is_empty() {
                    break;
                }

                // readline.add_history_entry(line.to_owned());
                lines.push(line.to_owned());
            },
            Ok(ReadlineEvent::Eof) => {
                break
            },
            Ok(ReadlineEvent::Interrupted) => {
                println!("^C");
                exit(1);
            },
            Err(error) => {
                error!("Unable to read in decrypted event source: {}", error.to_string());
                exit(1);
            },
        }
	}

    drop(_stdout);
    drop(readline);

    let event = match serde_json::from_str::<DecryptedEventSource>(lines.join("\n").as_str()) {
        Ok(event) => event,
        Err(error) => {
            error!("Unable to parse event: {}", error);
            exit(1);
        }
    };

    println!();

    debug!("Parsed event with {} content", event.content.msgtype.msgtype());

    if let Err(error) = match event.content.msgtype {
        MessageType::Audio(content) => process_media(content.source, Some(content.body), server).await,
        MessageType::File(content) => process_media(content.source, Some(content.body), server).await,
        MessageType::Image(content) => process_media(content.source, Some(content.body), server).await,
        MessageType::Video(content) => process_media(content.source, Some(content.body), server).await,
        _ => {
            error!("Event content is not known to contain media, exiting");
            exit(1);
        },
    } {
        error!("Unable to process media: {}", error);
    }

    Ok(())
}