joeyahines/FrenBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Removed imgur support

Browse Source 
+ Added my own basic image handling
+ Probably not perfect but works good enough in testing
+ clippy + fmt
This commit is contained in:
Joey Hines 2023-06-04 14:10:01 -06:00
parent 14032692b4
commit 7847bd38b9
Signed by: joeyahines
GPG Key ID: 995E531F7A569DDB
14 changed files with 339 additions and 257 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

126
src/album_manager/mod.rs Normal file
View File

@ -0,0 +1,126 @@
use crate::error::Error;
use j_db::database::Database;
use j_db::error::JDbError;
use j_db::model::JdbModel;
use rand::prelude::SliceRandom;
use reqwest::Url;
use serde::{Deserialize, Serialize};
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::path::PathBuf;
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Image {
pub filename: String,
pub tags: Vec<String>,
pub hash: u64,
pub db_name: String,
}
impl Image {
pub async fn new(data: &[u8], filename: &str, base_dir: PathBuf, tags: Vec<String>) -> Self {
let mut hasher = DefaultHasher::new();
data.hash(&mut hasher);
filename.hash(&mut hasher);
let hash = hasher.finish();
let file = PathBuf::from(filename);
let ext = file.extension().unwrap().to_str().unwrap();
let db_name = format!("{}.{}", hash, ext);
let path = base_dir.join(db_name.clone());
tokio::fs::write(path, data).await.unwrap();
Self {
filename: filename.to_string(),
tags,
hash,
db_name,
}
}
pub fn full_path(&self, base_dir: PathBuf) -> PathBuf {
base_dir.join(&self.db_name)
}
pub fn link(&self, base_link: &Url) -> Url {
base_link.join(&self.db_name).unwrap()
}
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Album {
id: Option<u64>,
pub album_name: String,
pub images: Vec<Image>,
pub aliases: Vec<String>,
}
impl Album {
pub fn find_album_by_name_or_alias(
db: &Database,
name: &str,
) -> Result<Option<Album>, JDbError> {
Ok(db
.filter(|_, album: &Album| {
album.album_name.eq_ignore_ascii_case(name)
|| album.aliases.contains(&name.to_string())
})?
.next())
}
pub fn get_random_image(
db: &Database,
album_name: &str,
tags: Vec<&str>,
) -> Result<Option<Image>, Error> {
let mut rng = rand::thread_rng();
let album = match Self::find_album_by_name_or_alias(db, album_name)? {
None => return Err(Error::NoAlbumFound),
Some(a) => a,
};
let images: Vec<Image> = if tags.is_empty() {
album.images
} else {
album
.images
.iter()
.filter(|img| {
for tag in &tags {
if img.tags.contains(&tag.to_lowercase()) {
return true;
}
}
false
})
.cloned()
.collect()
};
Ok(images.choose(&mut rng).cloned())
}
pub fn add_album(db: &Database, album_name: &str) -> Result<Self, JDbError> {
db.insert(Album {
id: None,
album_name: album_name.to_string(),
images: vec![],
aliases: vec![],
})
}
}
impl JdbModel for Album {
fn id(&self) -> Option<u64> {
self.id
}
fn set_id(&mut self, id: u64) {
self.id = Some(id)
}
fn tree() -> String {
"img_album".to_string()
}
}

67
src/config.rs
View File

@ -1,15 +1,12 @@
use crate::error::Error;
use crate::imgur;
use crate::imgur::Image;
use crate::models::album::AlbumConfig;
use crate::migrations::do_migration;
use config::{Config, File};
use j_db::database::Database;
use rand::prelude::SliceRandom;
use reqwest::Url;
use serde::{Deserialize, Serialize};
use serenity::model::id::ChannelId;
use serenity::model::prelude::{GuildId, UserId};
use serenity::prelude::TypeMapKey;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::sync::Arc;
@ -26,7 +23,8 @@ pub struct Args {
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct BotConfig {
pub bot_token: String,
pub imgur_client_id: String,
pub img_path: PathBuf,
pub base_url: Url,
pub story_path: PathBuf,
pub voice_path: PathBuf,
pub nft_path: PathBuf,
@ -50,66 +48,18 @@ impl BotConfig {
#[derive(Debug)]
pub struct BotState {
pub accepted_nsfw: Option<UserId>,
pub albums: HashMap<String, Vec<Image>>,
pub bad_apple_running: bool,
pub speak_lock: Mutex<()>,
}
impl BotState {
pub async fn new(cfg: &BotConfig, db: &Database) -> Result<Self, Error> {
let mut albums: HashMap<String, Vec<Image>> = HashMap::new();
for album in db.filter(|_, _: &AlbumConfig| true)? {
let imgur_album = imgur::get_album_images(&cfg.imgur_client_id, &album.album_id).await;
match imgur_album {
Ok(imgur) => {
albums.insert(album.name.clone(), imgur);
}
Err(err) => {
println!("Unable to parse album '{}'. Error: {}", album.name, err);
}
}
}
pub async fn new() -> Result<Self, Error> {
Ok(Self {
accepted_nsfw: None,
albums,
bad_apple_running: false,
speak_lock: Mutex::new(()),
})
}
pub fn get_image(&self, album_name: &str, tags: Vec<&str>) -> Result<Option<Image>, Error> {
let mut rng = rand::thread_rng();
let album = match self.albums.get(album_name) {
None => return Err(Error::NoAlbumFound),
Some(a) => a,
};
let album = if tags.is_empty() {
album.clone()
} else {
album
.iter()
.filter(|img| {
for tag in &tags {
if let Some(desc) = &img.description {
if desc.to_lowercase().contains(&tag.to_lowercase()) {
return true;
}
}
}
false
})
.cloned()
.collect()
};
Ok(album.choose(&mut rng).cloned())
}
}
pub struct GlobalData {
@ -122,9 +72,12 @@ pub struct GlobalData {
impl GlobalData {
pub async fn new(args: Args, cfg: BotConfig) -> Result<Self, Error> {
let db = Database::new(&cfg.db_path)?;
do_migration(&db);
Ok(Self {
args,
bot_state: BotState::new(&cfg, &db).await?,
bot_state: BotState::new().await?,
db,
cfg,
})
@ -134,7 +87,7 @@ impl GlobalData {
let cfg = BotConfig::new(&self.args.cfg_path)?;
self.cfg = cfg;
self.bot_state = BotState::new(&self.cfg, &self.db).await?;
self.bot_state = BotState::new().await?;
Ok(())
}

135
src/discord/album.rs
View File

@ -1,6 +1,6 @@
use crate::album_manager::{Album, Image};
use crate::discord::admin::is_admin;
use crate::error::Error;
use crate::imgur::get_album_images;
use crate::models::album::AlbumConfig;
use crate::{command, group, GlobalData};
use j_db::model::JdbModel;
use serenity::client::Context;
@ -8,63 +8,105 @@ use serenity::framework::standard::{Args, CommandResult};
use serenity::model::channel::Message;
#[group]
#[commands(add_album, remove_album, list_albums)]
pub struct Album;
#[commands(add_image, remove_album, list_albums, import_from_file)]
pub struct AlbumCmd;
#[command]
#[only_in(guilds)]
#[min_args(2)]
#[max_args(2)]
#[description("Add an imgur album command.")]
#[usage("<name> <album link>")]
async fn add_album(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
#[min_args(1)]
#[description("Add an image to an album")]
#[usage("<album_name>")]
async fn add_image(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let album_name = args.parse::<String>()?;
args.advance();
let album_id = args.parse::<String>()?;
let album_id: String = if album_id.contains("imgur") {
let parts: Vec<&str> = album_id.split("/a/").collect();
if parts.len() == 2 {
parts[1].to_string()
} else {
msg.reply(&ctx.http, "Invalid imgur album, check your link.")
.await?;
return Ok(());
}
} else {
album_id
};
let tags: Vec<String> = args.rest().split(',').map(|s| s.to_string()).collect();
let mut data = ctx.data.write().await;
let global_data = data.get_mut::<GlobalData>().unwrap();
let images = match get_album_images(&global_data.cfg.imgur_client_id, &album_id).await {
Ok(img) => img,
Err(e) => {
let album = Album::find_album_by_name_or_alias(&global_data.db, &album_name).unwrap();
let mut album = if let Some(album) = album {
album
} else {
Album::add_album(&global_data.db, &album_name).unwrap()
};
for attachment in &msg.attachments {
let data = attachment.download().await.unwrap();
let img = Image::new(
&data,
&attachment.filename,
global_data.cfg.img_path.clone(),
tags.clone(),
)
.await;
album.images.push(img);
}
global_data.db.insert(album).unwrap();
let plural = if msg.attachments.len() > 1 { "s" } else { "" };
msg.reply(
&ctx.http,
"Error adding album, check your link and try again",
format!("Image{} added to {}!", plural, album_name),
)
.await?;
println!("Error adding album: {}", e);
Ok(())
}
#[command]
#[only_in(guilds)]
#[min_args(1)]
#[description("Bulk add images")]
#[usage("<path>")]
async fn import_from_file(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
let album_name = args.parse::<String>()?;
args.advance();
let album_path = args.parse::<String>()?;
let album_path = album_path.replace('"', "");
let mut data = ctx.data.write().await;
let global_data = data.get_mut::<GlobalData>().unwrap();
if !is_admin(&msg.author.id, &global_data.cfg) {
return Ok(());
}
let album = Album::find_album_by_name_or_alias(&global_data.db, &album_name).unwrap();
let mut album = if let Some(album) = album {
album
} else {
Album::add_album(&global_data.db, &album_name).unwrap()
};
global_data
.db
.insert(AlbumConfig::new(&album_name, &album_id))?;
let mut dir = tokio::fs::read_dir(album_path).await.unwrap();
global_data
.bot_state
.albums
.insert(album_name.clone(), images);
while let Some(entry) = dir.next_entry().await? {
let data = tokio::fs::read(entry.path()).await?;
let img = Image::new(
&data,
entry.path().file_name().unwrap().to_str().unwrap(),
global_data.cfg.img_path.clone(),
vec![],
)
.await;
album.images.push(img);
}
msg.reply(&ctx.http, format!("{} album added!", album_name))
global_data.db.insert(album).unwrap();
let plural = if msg.attachments.len() > 1 { "s" } else { "" };
msg.reply(
&ctx.http,
format!("Image{} added to {}!", plural, album_name),
)
.await?;
Ok(())
@ -82,7 +124,7 @@ async fn remove_album(ctx: &Context, msg: &Message, args: Args) -> CommandResult
let global_data = data.get_mut::<GlobalData>().unwrap();
let album = AlbumConfig::get_album_by_name(&global_data.db, &album_name)?;
let album = Album::find_album_by_name_or_alias(&global_data.db, &album_name).unwrap();
match album {
None => {
@ -93,9 +135,13 @@ async fn remove_album(ctx: &Context, msg: &Message, args: Args) -> CommandResult
.await?;
}
Some(album) => {
global_data.bot_state.albums.remove(&album_name);
for img in &album.images {
tokio::fs::remove_file(img.full_path(global_data.cfg.img_path.clone()))
.await
.unwrap();
}
global_data.db.remove::<AlbumConfig>(album.id().unwrap())?;
global_data.db.remove::<Album>(album.id().unwrap())?;
msg.reply(&ctx.http, format!("{} album removed!", album_name))
.await?;
@ -116,8 +162,8 @@ async fn list_albums(ctx: &Context, msg: &Message, _args: Args) -> CommandResult
let album_names: Vec<String> = global_data
.db
.filter(|_, _: &AlbumConfig| true)?
.map(|album| album.name)
.filter(|_, _: &Album| true)?
.map(|album| album.album_name)
.collect();
if album_names.is_empty() {
@ -143,7 +189,7 @@ pub async fn parse_album(
let data = ctx.data.read().await;
let global_data = data.get::<GlobalData>().unwrap();
let img = match global_data.bot_state.get_image(album_name, tags) {
let img = match Album::get_random_image(&global_data.db, album_name, tags) {
Ok(img) => img,
Err(err) => {
return match err {
@ -154,7 +200,8 @@ pub async fn parse_album(
};
if let Some(img) = img {
msg.reply(&ctx.http, img.link).await?;
msg.reply(&ctx.http, img.link(&global_data.cfg.base_url))
.await?;
} else {
msg.reply(&ctx.http, "No image found :(").await?;
}

11
src/discord/joke.rs
View File

@ -1,3 +1,4 @@
use crate::album_manager::Album;
use crate::error::Error;
use crate::models::insult_compliment::{RandomResponseTemplate, ResponseType};
use crate::models::random::RandomConfig;
@ -55,11 +56,15 @@ impl RandomCtx {
pub fn new(user_name: &str, global_data: &GlobalData) -> Result<Self, Error> {
let mut random_image: HashMap<String, String> = HashMap::new();
for (album_name, images) in &global_data.bot_state.albums {
let image = images.choose(&mut thread_rng());
let albums: Vec<Album> = global_data.db.filter(|_, _album: &Album| true)?.collect();
for album in albums {
let image = album.images.choose(&mut thread_rng());
if let Some(image) = image {
random_image.insert(album_name.clone(), image.link.clone());
random_image.insert(
album.album_name.clone(),
image.link(&global_data.cfg.base_url).to_string(),
);
}
}

9
src/discord/motivate.rs
View File

@ -16,7 +16,7 @@ use std::borrow::Cow;
pub struct Motivate;
pub async fn create_motivation_image(motivation: Motivation) -> Result<Vec<u8>, CommandError> {
let motivation_image_blob = reqwest::get(&motivation.image_url).await?.bytes().await?;
let motivation_image_blob = tokio::fs::read(motivation.image_path).await?;
let text = format!("{} {}", motivation.action, motivation.goal);
@ -64,11 +64,8 @@ async fn motivation(ctx: &Context, msg: &Message, _args: Args) -> CommandResult
let data = ctx.data.read().await;
let global_data = data.get::<GlobalData>().unwrap();
let motivation = MotivationConfig::generate_motivation(
&global_data.db,
&global_data.cfg.imgur_client_id,
"white",
)
let motivation =
MotivationConfig::generate_motivation(&global_data.db, &global_data.cfg.img_path, "white")
.await?;
let image = create_motivation_image(motivation).await?;

2
src/discord/shop.rs
View File

@ -312,7 +312,7 @@ pub async fn restock_shop(ctx: &Context) -> Result<(), CommandError> {
let nft_motivation = MotivationConfig::generate_motivation(
&global_data.db,
&global_data.cfg.imgur_client_id,
&global_data.cfg.img_path,
"gold",
)
.await

9
src/error.rs
View File

@ -1,4 +1,3 @@
use crate::imgur::ImgurError;
use crate::user;
use serde::ser::StdError;
use std::fmt::{Display, Formatter};
@ -8,7 +7,6 @@ use std::fmt::{Display, Formatter};
#[allow(dead_code)]
pub enum Error {
ConfigError(config::ConfigError),
ImgurError(ImgurError),
SerenityError(serenity::Error),
TeraError(tera::Error),
NoAlbumFound,
@ -24,12 +22,6 @@ impl From<config::ConfigError> for Error {
}
}
impl From<ImgurError> for Error {
fn from(e: ImgurError) -> Self {
Self::ImgurError(e)
}
}
impl From<serenity::Error> for Error {
fn from(err: serenity::Error) -> Self {
Self::SerenityError(err)
@ -58,7 +50,6 @@ impl Display for Error {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Error::ConfigError(e) => write!(f, "Config error: {}", e),
Error::ImgurError(e) => write!(f, "Imgur error: {}", e),
Error::SerenityError(e) => write!(f, "Discord error: {}", e),
Error::TeraError(e) => write!(f, "Tera error: {}", e),
Error::NoAlbumFound => write!(f, "No album found"),

75
src/imgur/mod.rs
View File

@ -1,75 +0,0 @@
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::fmt::{Display, Formatter};
#[derive(Debug)]
pub enum ImgurError {
ReqwestError(reqwest::Error),
ImgurRequestError(String),
}
impl From<reqwest::Error> for ImgurError {
fn from(e: reqwest::Error) -> Self {
Self::ReqwestError(e)
}
}
impl Display for ImgurError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let msg = match self {
ImgurError::ReqwestError(err) => format!("Reqwest error: {}", err),
ImgurError::ImgurRequestError(msg) => format!("Imgur request error: {}", msg),
};
write!(f, "{}", msg)
}
}
impl std::error::Error for ImgurError {}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct AlbumData {
images: Option<Vec<Image>>,
error: Option<String>,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct AlbumResponse {
data: AlbumData,
success: bool,
status: i32,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Image {
pub id: String,
pub title: Option<String>,
pub description: Option<String>,
#[serde(rename = "type")]
pub img_type: String,
pub animated: bool,
pub width: i32,
pub height: i32,
pub size: i32,
pub link: String,
}
pub async fn get_album_images(client_id: &str, album_hash: &str) -> Result<Vec<Image>, ImgurError> {
let client = Client::new();
let res = client
.get(format!("https://api.imgur.com/3/album/{}", album_hash))
.header("Authorization", format!("Client-ID {}", client_id))
.send()
.await?;
let album_response: AlbumResponse = res.json().await?;
if album_response.success {
Ok(album_response.data.images.unwrap())
} else {
Err(ImgurError::ImgurRequestError(
album_response.data.error.unwrap(),
))
}
}

5
src/main.rs
View File

@ -1,9 +1,10 @@
mod album_manager;
mod api;
mod config;
mod discord;
mod error;
mod imgur;
mod inventory;
mod migrations;
mod models;
mod user;
@ -49,7 +50,7 @@ async fn main() {
let framework = StandardFramework::new()
.configure(|c| c.with_whitespace(true).prefix("!").ignore_bots(true))
.group(&discord::color::COLOR_GROUP)
.group(&discord::album::ALBUM_GROUP)
.group(&discord::album::ALBUMCMD_GROUP)
.group(&discord::celeryman::CELERYMAN_GROUP)
.group(&discord::joke::JOKE_GROUP)
.group(&discord::admin::ADMIN_GROUP)

43
src/migrations/migration2_remove_imgur.rs Normal file
View File

@ -0,0 +1,43 @@
use crate::models::motivation::MotivationConfig;
use j_db::database::Database;
use j_db::migration::Migration;
use j_db::model::JdbModel;
use json::JsonValue;
pub struct Migration2RemoveImgur {}
impl Migration for Migration2RemoveImgur {
fn up(&self, db: &Database) -> j_db::error::Result<()> {
db.db.open_tree("Album").unwrap().clear()?;
let mot_tree = db.db.open_tree(MotivationConfig::tree())?;
let (id, config) = mot_tree.iter().next().unwrap()?;
let mut config = json::parse(std::str::from_utf8(&config).unwrap()).unwrap();
config["album"] = JsonValue::Array(Vec::new());
mot_tree.insert(id, config.to_string().as_bytes())?;
Ok(())
}
fn down(&self, db: &Database) -> j_db::error::Result<()> {
let mot_tree = db.db.open_tree(MotivationConfig::tree())?;
let (id, config) = mot_tree.iter().next().unwrap()?;
let mut config = json::parse(std::str::from_utf8(&config).unwrap()).unwrap();
config["album"] = JsonValue::Array(Vec::new());
mot_tree.insert(id, config.to_string().as_bytes())?;
Ok(())
}
fn version(&self) -> u64 {
2
}
}

26
src/migrations/mod.rs Normal file
View File

@ -0,0 +1,26 @@
use crate::migrations::migration2_remove_imgur::Migration2RemoveImgur;
use j_db::database::Database;
use j_db::migration;
use j_db::migration::Direction;
mod migration2_remove_imgur;
const CURRENT_DB_VERSION: u64 = 2;
#[allow(clippy::single_match)]
pub fn do_migration(db: &Database) {
let version = db.version().unwrap_or(0);
for ver in version + 1..=CURRENT_DB_VERSION {
match ver {
2 => {
migration::do_migration::<Migration2RemoveImgur>(
db,
Migration2RemoveImgur {},
Direction::Up,
)
.unwrap();
}
_ => {}
}
}
}

41
src/models/album.rs
View File

@ -1,41 +0,0 @@
use crate::error::Error;
use j_db::database::Database;
use j_db::model::JdbModel;
use serde::{Deserialize, Serialize};
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct AlbumConfig {
id: Option<u64>,
pub name: String,
pub album_id: String,
}
impl JdbModel for AlbumConfig {
fn id(&self) -> Option<u64> {
self.id
}
fn set_id(&mut self, id: u64) {
self.id = Some(id)
}
fn tree() -> String {
"Album".to_string()
}
}
impl AlbumConfig {
pub fn new(name: &str, album_id: &str) -> AlbumConfig {
Self {
id: None,
name: name.to_string(),
album_id: album_id.to_string(),
}
}
pub fn get_album_by_name(db: &Database, name: &str) -> Result<Option<AlbumConfig>, Error> {
Ok(db
.filter(|_, album_config: &AlbumConfig| album_config.name.eq_ignore_ascii_case(name))?
.next())
}
}

1
src/models/mod.rs
View File

@ -1,4 +1,3 @@
pub mod album;
pub mod api_key;
pub mod birthday;
pub mod insult_compliment;

26
src/models/motivation.rs
View File

@ -1,15 +1,16 @@
use crate::album_manager::Album;
use crate::error::Error;
use crate::imgur::get_album_images;
use j_db::database::Database;
use j_db::model::JdbModel;
use rand::prelude::{IteratorRandom, SliceRandom};
use rand::thread_rng;
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct MotivationConfig {
id: Option<u64>,
pub album: Vec<String>,
pub album: Vec<u64>,
pub action: Vec<String>,
pub goal: Vec<String>,
}
@ -47,7 +48,10 @@ impl MotivationConfig {
pub fn add_album(db: &Database, album: &str) -> Result<(), Error> {
let mut motivation = Self::get_motivation_config(db)?;
motivation.album.push(album.to_string());
let album = Album::find_album_by_name_or_alias(db, album)
.unwrap()
.unwrap();
motivation.album.push(album.id().unwrap());
db.insert(motivation)?;
@ -74,7 +78,7 @@ impl MotivationConfig {
pub async fn generate_motivation(
db: &Database,
imgur_client_id: &str,
base_path: &Path,
border_color: &str,
) -> Result<Motivation, Error> {
let motivation = Self::get_motivation_config(db)?;
@ -82,18 +86,24 @@ impl MotivationConfig {
let mut images = Vec::new();
for album in motivation.album {
let mut album_images = get_album_images(imgur_client_id, &album).await?;
let album = db.get::<Album>(album);
if let Ok(album) = album {
let mut album_images = album.images.clone();
images.append(&mut album_images);
}
}
let image = images.choose(&mut thread_rng()).unwrap();
let image = images
.choose(&mut thread_rng())
.ok_or(Error::NoAlbumFound)?;
let action = motivation.action.iter().choose(&mut thread_rng()).unwrap();
let goal = motivation.goal.iter().choose(&mut thread_rng()).unwrap();
Ok(Motivation {
image_url: image.link.clone(),
image_path: image.full_path(base_path.to_path_buf()),
action: action.to_string(),
goal: goal.to_string(),
border_color: border_color.to_string(),
@ -103,7 +113,7 @@ impl MotivationConfig {
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct Motivation {
pub image_url: String,
pub image_path: PathBuf,
pub action: String,
pub goal: String,
pub border_color: String,