~zethra/mx-emote-thef

f90c6a5889a17a8a8532e72a81a1796063d9efb1 — Sashanoraa 1 year, 6 months ago 04c49d7 main
Tool is working now
5 files changed, 295 insertions(+), 125 deletions(-)

M .gitignore
M Cargo.lock
M Cargo.toml
A src/cli.rs
M src/main.rs
M .gitignore => .gitignore +2 -0
@@ 2,3 2,5 @@
.direnv
.envrc
.env
*.json
*.txt

M Cargo.lock => Cargo.lock +152 -1
@@ 23,6 23,55 @@ dependencies = [
]

[[package]]
name = "anstream"
version = "0.3.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ca84f3628370c59db74ee214b3263d58f9aadd9b4fe7e711fd87dc452b7f163"
dependencies = [
 "anstyle",
 "anstyle-parse",
 "anstyle-query",
 "anstyle-wincon",
 "colorchoice",
 "is-terminal",
 "utf8parse",
]

[[package]]
name = "anstyle"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "41ed9a86bf92ae6580e0a31281f65a1b1d867c0cc68d5346e2ae128dddfa6a7d"

[[package]]
name = "anstyle-parse"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e765fd216e48e067936442276d1d57399e37bce53c264d6fefbe298080cb57ee"
dependencies = [
 "utf8parse",
]

[[package]]
name = "anstyle-query"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5ca11d4be1bab0c8bc8734a9aa7bf4ee8316d462a08c6ac5052f888fef5b494b"
dependencies = [
 "windows-sys 0.48.0",
]

[[package]]
name = "anstyle-wincon"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "180abfa45703aebe0093f79badacc01b8fd4ea2e35118747e5811127f926e188"
dependencies = [
 "anstyle",
 "windows-sys 0.48.0",
]

[[package]]
name = "anyhow"
version = "1.0.71"
source = "registry+https://github.com/rust-lang/crates.io-index"


@@ 142,6 191,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9b1ce199063694f33ffb7dd4e0ee620741495c32833cde5aa08f02a0bf96f0c8"

[[package]]
name = "byteorder"
version = "1.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610"

[[package]]
name = "bytes"
version = "1.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"


@@ 154,12 209,65 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "50d30906286121d95be3d479533b458f87493b30a4b5f79a607db8f5d11aa91f"

[[package]]
name = "cfb"
version = "0.7.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d38f2da7a0a2c4ccf0065be06397cc26a81f4e528be095826eee9d4adbb8c60f"
dependencies = [
 "byteorder",
 "fnv",
 "uuid 1.3.2",
]

[[package]]
name = "cfg-if"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"

[[package]]
name = "clap"
version = "4.2.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "34d21f9bf1b425d2968943631ec91202fe5e837264063503708b83013f8fc938"
dependencies = [
 "clap_builder",
 "clap_derive",
 "once_cell",
]

[[package]]
name = "clap_builder"
version = "4.2.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "914c8c79fb560f238ef6429439a30023c862f7a28e688c58f7203f12b29970bd"
dependencies = [
 "anstream",
 "anstyle",
 "bitflags",
 "clap_lex",
 "strsim",
]

[[package]]
name = "clap_derive"
version = "4.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3f9644cd56d6b87dbe899ef8b053e331c0637664e9e21a33dfcdc36093f5c5c4"
dependencies = [
 "heck",
 "proc-macro2 1.0.56",
 "quote 1.0.26",
 "syn 2.0.15",
]

[[package]]
name = "clap_lex"
version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8a2dd5a6fe8c6e3502f568a6353e5273bbb15193ad9a89e457b9970798efbea1"

[[package]]
name = "cloudabi"
version = "0.0.3"
source = "registry+https://github.com/rust-lang/crates.io-index"


@@ 169,6 277,12 @@ dependencies = [
]

[[package]]
name = "colorchoice"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7"

[[package]]
name = "core-foundation"
version = "0.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"


@@ 511,6 625,12 @@ dependencies = [
]

[[package]]
name = "heck"
version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8"

[[package]]
name = "hermit-abi"
version = "0.2.6"
source = "registry+https://github.com/rust-lang/crates.io-index"


@@ 624,6 744,15 @@ dependencies = [
]

[[package]]
name = "infer"
version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f551f8c3a39f68f986517db0d1759de85881894fdc7db798bd2a9df9cb04b7fc"
dependencies = [
 "cfb",
]

[[package]]
name = "instant"
version = "0.1.12"
source = "registry+https://github.com/rust-lang/crates.io-index"


@@ 650,6 779,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "12b6ee2129af8d4fb011108c73d99a1b83a85977f23b82460c0ae2e25bb4b57f"

[[package]]
name = "is-terminal"
version = "0.4.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "adcf93614601c8129ddf72e2d5633df827ba6551541c6d8c59520a371475be1f"
dependencies = [
 "hermit-abi 0.3.1",
 "io-lifetimes",
 "rustix",
 "windows-sys 0.48.0",
]

[[package]]
name = "itoa"
version = "1.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"


@@ 834,13 975,17 @@ dependencies = [
]

[[package]]
name = "mx-emote-thef"
name = "mx-emote-replacer"
version = "0.1.0"
dependencies = [
 "anyhow",
 "clap",
 "dotenv",
 "infer",
 "matrix-sdk",
 "mime",
 "serde",
 "serde_json",
 "tokio",
]



@@ 1744,6 1889,12 @@ dependencies = [
]

[[package]]
name = "utf8parse"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "711b9620af191e0cdc7468a8d14e709c3dcdb115b36f838e601583af800a370a"

[[package]]
name = "uuid"
version = "0.7.4"
source = "registry+https://github.com/rust-lang/crates.io-index"

M Cargo.toml => Cargo.toml +6 -2
@@ 1,5 1,5 @@
[package]
name = "mx-emote-thef"
name = "mx-emote-replacer"
version = "0.1.0"
edition = "2021"



@@ 7,8 7,12 @@ edition = "2021"

[dependencies]
anyhow = "1.0.64" # Error handling
clap = { version = "4.2.7", features = ["derive"] }
dotenv = "0.15.0" # Load env vars for development
infer = "0.13.0"
mime = "0.3.17"
serde = { version = "1.0.145", features = ["derive"] }
serde_json = "1.0.96"

[dependencies.matrix-sdk]
version = "0.6.0"


@@ 17,4 21,4 @@ features = ["anyhow", "native-tls"]

[dependencies.tokio]
version = "1.21.0"
features = ["macros", "rt-multi-thread", "sync"] 
features = ["macros", "rt-multi-thread", "sync", "io-std"]

A src/cli.rs => src/cli.rs +20 -0
@@ 0,0 1,20 @@
use std::path::PathBuf;

use clap::Parser;
use matrix_sdk::ruma::OwnedUserId;

#[derive(Parser)]
#[command(author, version, about, long_about = None)]
pub struct Cli {
    /// If a media file can't be gotten from the account's HS, try these too
    pub matrix_servers: Vec<String>,
    /// Matrix user account to use for uploads
    #[arg(short = 'u', long)]
    pub maxtix_user: OwnedUserId,
    /// File to read session token from
    #[arg(short, long)]
    pub token_file: PathBuf,
    /// Check if images can be gotten without uploading new ones
    #[arg(long)]
    pub test: bool,
}

M src/main.rs => src/main.rs +115 -122
@@ 1,158 1,151 @@
mod cli;

use std::collections::HashMap;
use std::path::{Path, PathBuf};

use anyhow::Context;
use matrix_sdk::config::SyncSettings;
use matrix_sdk::media::MediaRequest;
use anyhow::{bail, Context};
use clap::Parser;
use matrix_sdk::media::{MediaFormat, MediaRequest};
use matrix_sdk::ruma::events::room::MediaSource;
use matrix_sdk::ruma::events::{EmptyStateKey, SyncStateEvent};
use matrix_sdk::ruma::exports::ruma_macros::EventContent;
use matrix_sdk::ruma::{OwnedRoomAliasId, OwnedUserId, RoomAliasId, OwnedMxcUri};
use matrix_sdk::Client;
use matrix_sdk::ruma::{device_id, MxcUri, OwnedMxcUri};
use matrix_sdk::{Client, Session};

use mime::Mime;
use serde::{Deserialize, Serialize};

const USAGE: &str = "usage: mx-emote-theif <output dir> <room alias>";
use serde_json::Value;
use tokio::io::{AsyncReadExt, AsyncWriteExt};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Get the path to save to from the args
    // Use args_os because paths don't have to UTF-8
    // Make this static by leaking memory (sneaky kitty coding)
    let save_dir: &'static PathBuf = Box::leak(Box::new(
        Path::new(
            &std::env::args_os()
                .nth(1)
                .ok_or_else(|| anyhow::anyhow!(USAGE))?,
        )
        .to_owned(),
    ));
    let room_alias: OwnedRoomAliasId = std::env::args_os()
        .nth(2)
        .ok_or_else(|| anyhow::anyhow!(USAGE))?
        .to_str()
        .ok_or_else(|| anyhow::anyhow!("<room alias> must be UTF-8"))?
        .try_into()
        .context("<room alias> must be a room alias")?;
    // Load .env if it's there
    if dotenv::dotenv().is_ok() {
        println!("Loaded .env");
    }
    // Get username and password from environment vars
    let username: OwnedUserId = std::env::var("MX_USER")
        .context("MX_USER must be set")?
        .try_into()
        .context("MX_USER must be valid user id")?;
    let password = std::env::var("MX_PASS").context("MX_PASS must be set")?;
    let cli = cli::Cli::parse();
    let hs_list: Vec<_> = cli
        .matrix_servers
        .iter()
        .map(|s| format!("https://{s}"))
        .collect();

    // Create output dir if it doesn't exist
    if !save_dir.exists() {
        tokio::fs::create_dir_all(&save_dir)
            .await
            .context("Error creating output dir")?;
    }
    let mut input_buf = Vec::new();
    tokio::io::stdin()
        .read_to_end(&mut input_buf)
        .await
        .context("Error reading event from stdin")?;

    let mut event: RoomEmotesEvent = serde_json::from_slice(&input_buf)
        .context("Input is not a valid im.ponies.room_emote event")?;

    // Make the matrix client and log in
    // // Make the matrix client and log in
    let client = Client::builder()
        .homeserver_url(format!("https://{}", username.server_name()))
        .homeserver_url(format!("https://{}", cli.maxtix_user.server_name()))
        .respect_login_well_known(true)
        .build()
        .await
        .context("Error setting up client")?;
    client
        .login_username(username.as_str(), &password)
        .initial_device_display_name("Room Avatar Theif")
        .send()
    let access_token = tokio::fs::read_to_string(cli.token_file)
        .await
        .context("Error login in")?;
    println!("Doing inital sync...");
        .context("Error reading token file")?
        .trim()
        .to_owned();
    let session = Session {
        access_token,
        refresh_token: None,
        user_id: cli.maxtix_user,
        device_id: device_id!("mx-emote-replacer").to_owned(),
    };
    client
        .sync_once(SyncSettings::default())
        .restore_login(session)
        .await
        .context("Initial sync failed")?;
    // Do big steal!
    if let Err(e) = steal_emotes(&client, save_dir, &room_alias).await {
        eprintln!("Error stealing images: {e:#}");
        .context("Error login in")?;

    for (name, image_info) in event.images.iter_mut() {
        eprintln!("Getting media for {name}");
        let image_data = match get_image(&image_info.url, &client, &hs_list).await {
            Ok(image_data) => image_data,
            Err(e) => {
                eprintln!("Couldn't get image {name}: {e}");
                if cli.test {
                    println!("{name}");
                }
                continue;
            }
        };
        if cli.test {
            continue;
        }
        let mime_type: Mime = match infer::get(&image_data) {
            Some(kind) => kind
                .mime_type()
                .parse()
                .context("infer libary returned an invalid mime type, this shouldn't happen")?,
            None => {
                eprintln!("{name} is not an image");
                continue;
            }
        };
        match client.media().upload(&mime_type, &image_data).await {
            Ok(resp) => image_info.url = resp.content_uri,
            Err(e) => eprintln!("Error uploading {name}: {e:#}"),
        }
    }
    if !cli.test {
        let output_event = serde_json::to_vec_pretty(&event).context("Error serialzing event")?;
        tokio::io::stdout().write_all(&output_event).await?;
    }
    // Log out even if stuff fails
    client.logout().await.context("Error logging out")?;

    Ok(())
}

async fn steal_emotes(
    client: &Client,
    save_dir: &'static Path,
    room_alias: &RoomAliasId,
) -> anyhow::Result<()> {
    println!("Starting...");
    let room_id = client
        .resolve_room_alias(room_alias)
        .await
        .context("Error resolving room alias")?
        .room_id;
    let room = client
        .get_joined_room(&room_id)
        .ok_or_else(|| anyhow::anyhow!("You must be in the selected room"))?;
    let SyncStateEvent::Original(emote_event) = room
        .get_state_event_static::<RoomEmotesContent>()
        .await
        .context("Error getting room emotes event")?
        .ok_or_else(|| anyhow::anyhow!("Room does not have an emotes event"))?
        .deserialize().context("Error parsing emote event")? else {
        anyhow::bail!("Room only has a redacted emote event");
    };
    for (name, join_handle) in emote_event.content.images.into_iter().map(|(name, image)| {
        let client = client.clone();
        (
            name.clone(),
            tokio::spawn(async move {
                println!("Stealing {name}");
                let image_data = client
                    .media()
                    .get_media_content(
                        &MediaRequest {
                            source: MediaSource::Plain(image.url),
                            format: matrix_sdk::media::MediaFormat::File,
                        },
                        true,
                    )
                    .await
                    .context("Error getting image")?;
                tokio::fs::write(save_dir.join(&name), image_data)
                    .await
                    .with_context(|| anyhow::anyhow!("Error writing image {name}"))?;
                anyhow::Result::<()>::Ok(())
            }),
async fn get_image(url: &MxcUri, client: &Client, hs_list: &[String]) -> anyhow::Result<Vec<u8>> {
    match client
        .media()
        .get_media_content(
            &MediaRequest {
                source: MediaSource::Plain(url.to_owned()),
                format: MediaFormat::File,
            },
            false,
        )
    }) {
        // Wait for all the processes to finish before returning
        match join_handle.await {
            Ok(ret) => {
                if let Err(e) = ret {
                    eprintln!("Error getting image for emote {name}: {e:#}")
        .await
    {
        Ok(image_file) => Ok(image_file),
        Err(_) => {
            for server in hs_list {
                if let Ok(image_data) = get_image_from_server(url, server).await {
                    return Ok(image_data);
                }
            }
            Err(e) => eprintln!("Erro waiting on process for emote {name}: {e:#}"),
            bail!("Couldn't get image from any home server");
        }
    }
    Ok(())
}

#[derive(Clone, Debug, Deserialize, Serialize, EventContent)]
#[ruma_event(type = "im.ponies.room_emotes", kind = State, state_key_type = EmptyStateKey)]
pub struct RoomEmotesContent {
    images: HashMap<String, EmoteImage>,
    pack: EmotePackInfo,
async fn get_image_from_server(url: &MxcUri, server: &str) -> anyhow::Result<Vec<u8>> {
    eprintln!("Falling back to server {server}");
    let client = Client::builder()
        .homeserver_url(server)
        .respect_login_well_known(true)
        .build()
        .await
        .with_context(|| format!("Error setting up client for server {server}"))?;
    Ok(client
        .media()
        .get_media_content(
            &MediaRequest {
                source: MediaSource::Plain(url.to_owned()),
                format: MediaFormat::File,
            },
            false,
        )
        .await?)
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct EmoteImage {
    url: OwnedMxcUri,
    usage: Option<Vec<String>>,
pub struct RoomEmotesEvent {
    images: HashMap<String, EmoteImage>,
    pack: Option<Value>,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct EmotePackInfo {
    display_name: Option<String>,
    usage: Option<Vec<String>>,
pub struct EmoteImage {
    url: OwnedMxcUri,
    usage: Option<Value>,
    info: Option<Value>,
}