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) { let mut media_url = "".to_string(); let mut media_ids: Vec = 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::(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 { 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> { 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> { // 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("=")); } }