mastodon-image-uploader-bot/src/main.rs

use async_std;
use log;
use mastodon_async::entities::visibility::Visibility;
use mastodon_async::helpers::{cli, toml as masto_toml};
use mastodon_async::prelude::*;
use reqwest;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::io::{Cursor, Write};
use std::process::exit;
use std::time::Duration;
use toml;

#[derive(Debug, Serialize)]
struct AccountUpdate {
  note: String,
}

const CONFIG_FILENAME: &str = "config.toml";

type DynResult<T> = Result<T, Box<dyn std::error::Error>>;

#[derive(Deserialize)]
struct Config {
  bot: Bot,
  files: Files,
  errors: Errors,
}

#[derive(Deserialize)]
struct Bot {
  name: String,
  instance: String,
  bio: String,
}

#[derive(Deserialize)]
struct Files {
  posted: String,
  urls: String,
  tempfile: String,
}

#[derive(Deserialize)]
struct Errors {
  maintainers: String,
  out_of_images: String,
  retry: u8,
}

#[tokio::main] // requires `features = ["mt"]
async fn main() -> DynResult<()> {
  stderrlog::new()
    .module(module_path!())
    .quiet(false)
    .verbosity(4) // Debug
    .timestamp(stderrlog::Timestamp::Second)
    .init()?;

  let config: Config = get_config();
  let mastodon = get_account(&config).await;

  match get_next_url(&config) {
    Ok(url) => match url {
      Some(url) => {
        let mut retry: u8 = 0;
        while let Err(err) = post_image(&mastodon, &url, &config, Visibility::Unlisted).await {
          log::warn!("Cannot post image, retry: {}, {}", retry, err);
          async_std::task::sleep(Duration::new(1, 0)).await;
          retry += 1;
          if retry >= config.errors.retry {
            log::error!("Max ammount of retries reached on post_image");
            log::info!("Reverting file changes");
            exit(1);
          }
        }
        set_url_as_posted(&config, &url)?;
        let mut retry: u8 = 0;
        while let Err(err) = update_bio(&mastodon, &config).await {
          log::warn!("Cannot update bio, retry: {}, {}", retry, err);
          async_std::task::sleep(Duration::new(1, 0)).await;
          retry += 1;
          if retry >= config.errors.retry {
            log::error!("Max ammount of retries reached on update bio");
            exit(1);
          }
        }
      }
      None => {
        let mut retry: u8 = 0;
        while let Err(err) = post(
          &mastodon,
          &format!(
            "{} {}",
            &config.errors.maintainers, &config.errors.out_of_images
          ),
          Visibility::Direct,
        )
        .await
        {
          log::warn!("Cannot post, retry: {}, {}", retry, err);
          async_std::task::sleep(Duration::new(1, 0)).await;
          retry += 1;
          if retry >= config.errors.retry {
            log::error!("Max ammount of retries reached on post");
            exit(1);
          }
        }
      }
    },
    Err(err) => {
      log::error!("Cannot get next image: {}", err);
      match post(
        &mastodon,
        &format!("{} {}", &config.errors.maintainers, &err.to_string()),
        Visibility::Direct,
      )
      .await
      {
        Ok(_) => {}
        Err(err) => {
          log::error!("Cannot post error message: {}", err);
          exit(1);
        }
      };
    }
  }
  Ok(())
}

fn get_config() -> Config {
  match parse_config(CONFIG_FILENAME) {
    Ok(config) => config,
    Err(err) => {
      log::error!("Config file parsing unsuccesful: {}", err);
      exit(1);
    }
  }
}

async fn get_account(config: &Config) -> Mastodon {
  if let Ok(data) = masto_toml::from_file("mastodon-data.toml") {
    Mastodon::from(data)
  } else {
    match register(config).await {
      Ok(account) => account,
      Err(err) => {
        log::error!("Api registation unsuccesful: {}", err);
        exit(1);
      }
    }
  }
}

/// Parses the given filename to a config struct
fn parse_config(filename: &str) -> DynResult<Config> {
  let toml_file = std::fs::read_to_string(filename)?; //.expect("No config file, consider getting the original one and modifing it");
  Ok(toml::from_str(&toml_file)?) //("Malformed config file, check the original one for reference")
}

fn get_next_url(config: &Config) -> DynResult<Option<String>> {
  let binding = std::fs::read_to_string(&config.files.posted)?; //.expect("Posted file not found");
  let posted = binding.lines().collect::<HashSet<&str>>();
  let binding = std::fs::read_to_string(&config.files.urls)?; //.expect("Urls file not found");
  let urls = binding.lines().collect::<HashSet<&str>>();
  let urls = urls.difference(&posted).collect::<Vec<_>>();
  if urls.is_empty() {
    Ok(None)
  } else {
    Ok(Some(urls[0].to_string().clone()))
  }
}

fn set_url_as_posted(config: &Config, url: &String) -> DynResult<()> {
  let mut file = std::fs::OpenOptions::new()
    .write(true)
    .append(true) // This is needed to append to file
    .open(&config.files.posted)?; //.expect("Cannot open posted file"); // Maybe we should retry just in case
  write!(file, "{}\n", url)?; //.expect("Cannot write to posted file"); // maybe we should retry tbh
  log::info!("Set url {} as posted", url);
  Ok(())
}

async fn post_image(
  account: &Mastodon,
  url: &String,
  config: &Config,
  visibility: Visibility,
) -> DynResult<Status> {
  fetch_url(url, &config.files.tempfile).await?; //.expect("Error fetching url");
  let attachment = account
    .media(&config.files.tempfile, Some(url.to_string()))
    .await?; //.expect("Attachment upload error");
  let attachment = account
    .wait_for_processing(attachment, Default::default())
    .await?; //.expect("Attachment processing error");
  let status = StatusBuilder::new()
    .media_ids(&[attachment.id])
    .visibility(visibility)
    .sensitive(true)
    .build()?; //.expect("Could not build status"); // we should retry
  let status = account.new_status(status).await?; //.expect("Error generating status"); // we should retry or delete last url in posted
  log::info!("Image status posted: {}", url);
  Ok(status)
}

async fn update_bio(account: &Mastodon, config: &Config) -> DynResult<()> {
  let binding = std::fs::read_to_string(&config.files.posted)?; //.expect("Posted file not found");
  let posted = binding.lines().collect::<HashSet<&str>>();
  let binding = std::fs::read_to_string(&config.files.urls)?; //.expect("Url file not found");
  let urls = binding.lines().collect::<HashSet<&str>>();

  let remaining = urls.difference(&posted).count();
  let client = reqwest::Client::new();

  let account_update = AccountUpdate {
    note: format!("{}\n\n{} new images remaining", config.bot.bio, remaining),
  };

  let response = client
    .patch(&format!(
      "{}/api/v1/accounts/update_credentials",
      config.bot.instance
    ))
    .bearer_auth(&account.data.token)
    .json(&account_update)
    .send()
    .await?;
  log::info!("Bio updated, response status: {}", response.status());
  Ok(())
}

async fn post(account: &Mastodon, msg: &str, visibility: Visibility) -> DynResult<Status> {
  let status = StatusBuilder::new()
    .visibility(visibility)
    .status(msg)
    .build()?; //.expect("Error building error status");
  let post = account.new_status(status).await?; //.expect("Error posting error status");
  log::info!("Text status posted: {}", msg);
  Ok(post)
}

async fn fetch_url(url: &String, file_name: &String) -> DynResult<()> {
  let response = reqwest::get(url);
  let mut file = std::fs::File::create(file_name)?;
  let mut content = Cursor::new(response.await?.bytes().await?);
  std::io::copy(&mut content, &mut file)?;
  Ok(())
}

async fn register(config: &Config) -> DynResult<Mastodon> {
  let registration = Registration::new(&config.bot.instance)
    .client_name(&config.bot.name)
    .scopes(Scopes::all())
    .build()
    .await?;
  let mastodon = cli::authenticate(registration).await?;

  // Save app data for using on the next run.
  masto_toml::to_file(&mastodon.data, "mastodon-data.toml")?;

  Ok(mastodon)
}

#[cfg(test)]
mod tests {

  use reqwest::StatusCode;

  use super::*;
  const TMPTESTDIR: &str = "/tmp/botimage";
  const TEST_URL: &str = "https://2.gravatar.com/avatar/be8eb8426d68e4beb50790647eda6f6b";

  #[tokio::test]
  async fn fetch_url_should_fetch() {
    fetch_url(&TEST_URL.to_string(), &TMPTESTDIR.to_string())
      .await
      .unwrap();
    std::fs::read(TMPTESTDIR).unwrap();
  }

  #[tokio::test]
  async fn post_should_post() {
    let client = reqwest::Client::new();
    let config = get_config();
    let account = get_account(&config).await;
    let msg = format!("Test!");

    let status = post(&account, &msg, Visibility::Direct).await.unwrap();
    let response = client
      .get(dbg!(format!(
        "{}/api/v1/statuses/{}",
        &config.bot.instance,
        &status.id.to_string()
      )))
      .bearer_auth(dbg!(&account.data.token))
      .send()
      .await
      .unwrap();
    account.delete_status(&status.id).await.unwrap();

    assert_eq!(response.status(), StatusCode::OK)
  }

  #[tokio::test]
  async fn post_image_works() {
    let client = reqwest::Client::new();
    let config = get_config();
    let account = get_account(&config).await;

    let status = post_image(&account, &TEST_URL.to_string(), &config, Visibility::Direct)
      .await
      .unwrap();

    let response = account.get_status(&status.id).await.unwrap();

    account.delete_status(&status.id).await.unwrap();
    let attachment = &response.media_attachments[0];

    let response = client
      .get(attachment.url.clone().unwrap())
      .bearer_auth(dbg!(&account.data.token))
      .send()
      .await
      .unwrap();

    assert_eq!(response.status(), StatusCode::OK)
  }

  #[test]
  fn set_as_posted_works() {
    let config = get_config();

    std::fs::write(&config.files.posted, "").unwrap();

    set_url_as_posted(&config, &TEST_URL.to_string()).unwrap();
    let file = std::fs::read_to_string(&config.files.posted).unwrap();
    let url = file.lines().next_back().unwrap();
    assert_eq!(url, TEST_URL)
  }

  #[test]
  fn get_next_url_works() {
    let config = get_config();
    // Reset file
    std::fs::write(&config.files.posted, "").unwrap();

    // Get test url
    let url = get_next_url(&config).unwrap().unwrap();
    assert_eq!(url, TEST_URL);

    set_url_as_posted(&config, &TEST_URL.to_string()).unwrap();
    let file = std::fs::read_to_string(&config.files.posted).unwrap();
    let url = file.lines().next_back().unwrap();
    assert_eq!(url, TEST_URL);

    // Test that now it does not get it
    let url = get_next_url(&config).unwrap();
    assert_eq!(url, None);
  }
}