scootaloo/src/util.rs

use crate::{twitter::get_tweet_media, ScootalooError};

use base64::encode;
use egg_mode::tweet::Tweet;
use futures::{stream, stream::StreamExt};
use log::{error, info, warn};
use megalodon::{
    entities::UploadMedia::{AsyncAttachment, Attachment},
    error,
    mastodon::Mastodon,
    megalodon::Megalodon,
};
use reqwest::Url;
use std::error::Error;
use tokio::{
    fs::{create_dir_all, remove_file, File},
    io::copy,
    time::{sleep, Duration},
};

/// Generate associative table between media ids and tweet extended entities
pub async fn generate_media_ids(
    tweet: &Tweet,
    cache_path: &str,
    mastodon: &Mastodon,
) -> (String, Vec<String>) {
    let mut media_url = "".to_string();
    let mut media_ids: Vec<String> = vec![];

    if let Some(m) = &tweet.extended_entities {
        info!("{} medias in tweet", m.media.len());

        let medias = m.media.clone();

        let mut stream = stream::iter(medias)
            .map(|media| {
                // attribute media url
                media_url = media.url.clone();

                // clone everything we need
                let cache_path = String::from(cache_path);
                let mastodon = mastodon.clone();

                tokio::task::spawn(async move {
                    info!("Start treating {}", media.media_url_https);
                    // get the tweet embedded media
                    let local_tweet_media_path = get_tweet_media(&media, &cache_path).await?;

                    // upload media to Mastodon
                    let mastodon_media = mastodon
                        .upload_media(local_tweet_media_path.to_owned(), None)
                        .await?
                        .json();
                    // at this point, we can safely erase the original file
                    // it doesn’t matter if we can’t remove, cache_media fn is idempotent
                    remove_file(&local_tweet_media_path).await.ok();

                    let id = match mastodon_media {
                        Attachment(m) => m.id,
                        AsyncAttachment(m) => wait_until_uploaded(&mastodon, &m.id).await?,
                    };

                    Ok::<String, ScootalooError>(id)
                })
            })
            .buffered(4); // there are max four medias per tweet and they need to be treated in
                          // order

        while let Some(result) = stream.next().await {
            match result {
                Ok(Ok(v)) => media_ids.push(v),
                Ok(Err(e)) => warn!("Cannot treat media: {}", e),
                Err(e) => error!("Something went wrong when joining the main thread: {}", e),
            }
        }
    } else {
        info!("No media in tweet");
    }

    // in case some media_ids slot remained empty due to errors, remove them
    media_ids.retain(|x| !x.is_empty());

    (media_url, media_ids)
}

/// Wait on uploaded medias when necessary
async fn wait_until_uploaded(client: &Mastodon, id: &str) -> Result<String, error::Error> {
    loop {
        sleep(Duration::from_secs(1)).await;
        let res = client.get_media(id.to_string()).await;
        return match res {
            Ok(res) => Ok(res.json.id),
            Err(err) => match err {
                error::Error::OwnError(ref own_err) => match own_err.kind {
                    error::Kind::HTTPPartialContentError => continue,
                    _ => Err(err),
                },
                _ => Err(err),
            },
        };
    }
}

/// Transforms the media into a base64 equivalent
pub async fn base64_media(u: &str) -> Result<String, Box<dyn Error>> {
    let mut response = reqwest::get(u).await?;

    let mut buffer = Vec::new();

    while let Some(chunk) = response.chunk().await? {
        copy(&mut &*chunk, &mut buffer).await?;
    }

    let content_type = response
        .headers()
        .get("content-type")
        .ok_or_else(|| ScootalooError::new(&format!("Cannot get media content type for {u}")))?
        .to_str()?;

    let encoded_f = encode(buffer);

    Ok(format!("data:{content_type};base64,{encoded_f}"))
}

/// Gets and caches Twitter Media inside the determined temp dir
pub async fn cache_media(u: &str, t: &str) -> Result<String, Box<dyn Error>> {
    // create dir
    create_dir_all(t).await?;

    // get file
    let mut response = reqwest::get(u).await?;

    // create local file
    let url = Url::parse(u)?;
    let dest_filename = url
        .path_segments()
        .ok_or_else(|| {
            ScootalooError::new(&format!(
                "Cannot determine the destination filename for {u}"
            ))
        })?
        .last()
        .ok_or_else(|| {
            ScootalooError::new(&format!(
                "Cannot determine the destination filename for {u}"
            ))
        })?;

    let dest_filepath = format!("{t}/{dest_filename}");

    let mut dest_file = File::create(&dest_filepath).await?;

    while let Some(chunk) = response.chunk().await? {
        copy(&mut &*chunk, &mut dest_file).await?;
    }

    Ok(dest_filepath)
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::{fs::remove_dir_all, path::Path};

    const TMP_DIR: &'static str = "/tmp/scootaloo_test";

    #[tokio::test]
    async fn test_cache_media() {
        let dest = cache_media(
            "https://forum.nintendojo.fr/styles/prosilver/theme/images/ndfr_casual.png",
            TMP_DIR,
        )
        .await
        .unwrap();

        assert!(Path::new(&dest).exists());

        remove_dir_all(TMP_DIR).unwrap();
    }

    #[tokio::test]
    async fn test_base64_media() {
        let img = base64_media(
            "https://forum.nintendojo.fr/styles/prosilver/theme/images/ndfr_casual.png",
        )
        .await
        .unwrap();

        assert!(img.starts_with("data:image/png;base64,"));

        assert!(img.ends_with("="));
    }
}