pikadick/commands/
deviantart.rs

1use crate::{
2    checks::ENABLED_CHECK,
3    client_data::{
4        CacheStatsBuilder,
5        CacheStatsProvider,
6    },
7    util::{
8        LoadingReaction,
9        TimedCache,
10        TimedCacheEntry,
11    },
12    ClientDataKey,
13    Database,
14};
15use anyhow::Context as _;
16use deviantart::Deviation;
17use rand::seq::IteratorRandom;
18use serenity::{
19    framework::standard::{
20        macros::command,
21        Args,
22        CommandResult,
23    },
24    model::prelude::*,
25    prelude::*,
26};
27use std::{
28    sync::Arc,
29    time::Instant,
30};
31use tracing::{
32    error,
33    info,
34};
35
36const DATA_STORE_NAME: &str = "deviantart";
37const COOKIE_KEY: &str = "cookie-store";
38
39/// A caching deviantart client
40#[derive(Clone, Debug)]
41pub struct DeviantartClient {
42    client: deviantart::Client,
43    search_cache: TimedCache<String, Vec<Deviation>>,
44}
45
46impl DeviantartClient {
47    /// Make a new [`DeviantartClient`].
48    pub async fn new(db: &Database) -> anyhow::Result<Self> {
49        use std::io::BufReader;
50
51        let client = deviantart::Client::new();
52
53        let cookie_data: Option<Vec<u8>> = db
54            .store_get(DATA_STORE_NAME, COOKIE_KEY)
55            .await
56            .context("failed to get cookie data")?;
57
58        match cookie_data {
59            Some(cookie_data) => {
60                client
61                    .load_json_cookies(BufReader::new(std::io::Cursor::new(cookie_data)))
62                    .await?;
63            }
64            None => {
65                info!("could not load cookie data");
66            }
67        }
68
69        Ok(DeviantartClient {
70            client,
71            search_cache: TimedCache::new(),
72        })
73    }
74
75    /// Signs in if necessary
76    pub async fn sign_in(
77        &self,
78        db: &Database,
79        username: &str,
80        password: &str,
81    ) -> anyhow::Result<()> {
82        if !self.client.is_logged_in_online().await? {
83            info!("re-signing in");
84            self.client.login(username, password).await?;
85
86            // Store the new cookies
87            let cookie_store = self.client.cookie_store.clone();
88            let cookie_data = tokio::task::spawn_blocking(move || {
89                let mut cookie_data = Vec::with_capacity(1_000_000); // 1 MB
90                cookie_store
91                    .lock()
92                    .expect("cookie store is poisoned")
93                    .save_json(&mut cookie_data)
94                    .map_err(deviantart::WrapBoxError)?;
95                anyhow::Result::<_>::Ok(cookie_data)
96            })
97            .await??;
98            db.store_put(DATA_STORE_NAME, COOKIE_KEY, cookie_data)
99                .await?;
100        }
101
102        Ok(())
103    }
104
105    /// Search for deviantart images with a cache.
106    pub async fn search(
107        &self,
108        db: &Database,
109        username: &str,
110        password: &str,
111        query: &str,
112    ) -> anyhow::Result<Arc<TimedCacheEntry<Vec<Deviation>>>> {
113        if let Some(entry) = self.search_cache.get_if_fresh(query) {
114            return Ok(entry);
115        }
116
117        let start = Instant::now();
118        self.sign_in(db, username, password)
119            .await
120            .context("failed to log in to deviantart")?;
121        let mut search_cursor = self.client.search(query, None);
122        search_cursor
123            .next_page()
124            .await
125            .context("failed to search")?;
126        let list = search_cursor
127            .take_current_deviations()
128            .expect("missing page")
129            .context("failed to process results")?;
130        let ret = self.search_cache.insert_and_get(String::from(query), list);
131
132        info!("searched deviantart in {:?}", start.elapsed());
133
134        Ok(ret)
135    }
136}
137
138impl CacheStatsProvider for DeviantartClient {
139    fn publish_cache_stats(&self, cache_stats_builder: &mut CacheStatsBuilder) {
140        cache_stats_builder.publish_stat(
141            "deviantart",
142            "search_cache",
143            self.search_cache.len() as f32,
144        );
145    }
146}
147
148#[command]
149#[description("Get art from deviantart")]
150#[usage("<query>")]
151#[example("sun")]
152#[min_args(1)]
153#[max_args(1)]
154#[checks(Enabled)]
155#[bucket("default")]
156async fn deviantart(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
157    let data_lock = ctx.data.read().await;
158    let client_data = data_lock
159        .get::<ClientDataKey>()
160        .expect("missing clientdata");
161    let client = client_data.deviantart_client.clone();
162    let db = client_data.db.clone();
163    let config = client_data.config.clone();
164    drop(data_lock);
165
166    let query = args.trimmed().quoted().current().expect("missing query");
167
168    info!("Searching for '{}' on deviantart", query);
169
170    let mut loading = LoadingReaction::new(ctx.http.clone(), msg);
171
172    match client
173        .search(
174            &db,
175            &config.deviantart.username,
176            &config.deviantart.password,
177            query,
178        )
179        .await
180    {
181        Ok(entry) => {
182            let data = entry.data();
183            let choice = data
184                .iter()
185                .filter_map(|deviation| {
186                    if deviation.is_image() {
187                        Some(
188                            deviation
189                                .get_image_download_url()
190                                .or_else(|| deviation.get_fullview_url()),
191                        )
192                    } else if deviation.is_film() {
193                        Some(deviation.get_best_video_url().cloned())
194                    } else {
195                        None
196                    }
197                })
198                .choose(&mut rand::thread_rng());
199
200            if let Some(choice) = choice {
201                if let Some(url) = choice {
202                    loading.send_ok();
203                    msg.channel_id.say(&ctx.http, url).await?;
204                } else {
205                    msg.channel_id
206                        .say(&ctx.http, "Missing url. This is probably a bug.")
207                        .await?;
208                    error!("DeviantArt deviation missing asset url: {:?}", choice);
209                }
210            } else {
211                msg.channel_id.say(&ctx.http, "No Results").await?;
212            }
213        }
214        Err(e) => {
215            msg.channel_id
216                .say(&ctx.http, format!("Failed to search '{}': {:?}", query, e))
217                .await?;
218
219            error!("Failed to search for {} on deviantart: {:?}", query, e);
220        }
221    }
222
223    client.search_cache.trim();
224
225    Ok(())
226}
pikadick/commands/deviantart.rs

pikadick/commands/
deviantart.rs
