Pull info from an attached .osu/.osz file (#18)

youmubot-osu/Cargo.toml

@@ -7,16 +7,18 @@ edition = "2018"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-serenity = "0.10"
-chrono = "0.4"
-reqwest = "0.11"
-serde = { version = "1.0", features = ["derive"] }
-bitflags = "1"
-lazy_static = "1"
-regex = "1"
-rosu-pp = "0.4"
-dashmap = "4"
 bincode = "1"
+bitflags = "1"
+chrono = "0.4"
+dashmap = "4"
+lazy_static = "1"
+osuparse = { git = "https://github.com/eltrufas/osuparse", rev = "ad8f6e5e7771e7cbaa2ec96c376558f9731139af" }
+regex = "1"
+reqwest = "0.11"
+rosu-pp = "0.4"
+serde = { version = "1.0", features = ["derive"] }
+serenity = "0.10"
+zip = "0.6.2"
 
 youmubot-db = { path = "../youmubot-db" }
 youmubot-db-sql = { path = "../youmubot-db-sql" }

youmubot-osu/src/discord/embeds.rs

@@ -1,9 +1,10 @@
 use super::BeatmapWithMode;
 use crate::{
     discord::oppai_cache::{Accuracy, BeatmapContent, BeatmapInfo, BeatmapInfoWithPP},
-    models::{Beatmap, Mode, Mods, Rank, Score, User},
+    models::{Beatmap, Difficulty, Mode, Mods, Rank, Score, User},
 };
 use serenity::{builder::CreateEmbed, utils::MessageBuilder};
+use std::time::Duration;
 use youmubot_prelude::*;
 
 /// Writes a number grouped in groups of 3.
@@ -55,42 +56,54 @@ fn beatmap_description(b: &Beatmap) -> String {
         .build()
 }
 
-pub fn beatmap_embed<'a>(
-    b: &'_ Beatmap,
+pub fn beatmap_offline_embed(
+    b: &'_ crate::discord::oppai_cache::BeatmapContent,
     m: Mode,
     mods: Mods,
-    info: Option<BeatmapInfoWithPP>,
-    c: &'a mut CreateEmbed,
-) -> &'a mut CreateEmbed {
-    let mod_str = if mods == Mods::NOMOD {
-        "".to_owned()
+) -> Result<Box<dyn FnOnce(&mut CreateEmbed) -> &mut CreateEmbed + Send + Sync>> {
+    let bm = b.content.clone();
+    let metadata = b.metadata.clone();
+    let (info, pp) = b.get_possible_pp_with(mods)?;
+
+    let total_length = if bm.hit_objects.len() >= 1 {
+        Duration::from_millis(
+            (bm.hit_objects.last().unwrap().end_time() - bm.hit_objects.first().unwrap().start_time)
+                as u64,
+        )
     } else {
-        format!(" {}", mods)
+        Duration::from_secs(0)
     };
-    let diff = b
-        .difficulty
-        .apply_mods(mods, info.map(|(v, _)| v.stars as f64));
-    c.title(
-        MessageBuilder::new()
-            .push_bold_safe(&b.artist)
-            .push(" - ")
-            .push_bold_safe(&b.title)
-            .push(" [")
-            .push_bold_safe(&b.difficulty_name)
-            .push("]")
-            .push(&mod_str)
-            .build(),
-    )
-    .author(|a| {
-        a.name(&b.creator)
-            .url(format!("https://osu.ppy.sh/users/{}", b.creator_id))
-            .icon_url(format!("https://a.ppy.sh/{}", b.creator_id))
-    })
-    .url(b.link())
-    .image(b.cover_url())
-    .color(0xffb6c1)
-    .fields(info.map(|(_, pp)| {
-        (
+
+    let diff = Difficulty {
+        stars: info.stars,
+        aim: None,   // TODO: this is currently unused
+        speed: None, // TODO: this is currently unused
+        cs: bm.cs as f64,
+        od: bm.od as f64,
+        ar: bm.ar as f64,
+        hp: bm.hp as f64,
+        count_normal: bm.n_circles as u64,
+        count_slider: bm.n_sliders as u64,
+        count_spinner: bm.n_spinners as u64,
+        max_combo: Some(info.max_combo as u64),
+        bpm: bm.bpm(),
+        drain_length: total_length, // It's hard to calculate so maybe just skip...
+        total_length,
+    }
+    .apply_mods(mods, Some(info.stars));
+    Ok(Box::new(move |c: &mut CreateEmbed| {
+        c.title(beatmap_title(
+            &metadata.artist,
+            &metadata.title,
+            &metadata.version,
+            mods,
+        ))
+        .author(|a| {
+            a.name(&metadata.creator)
+                .url(format!("https://osu.ppy.sh/users/{}", metadata.creator))
+        })
+        .color(0xffb6c1)
+        .field(
             "Calculated pp",
             format!(
                 "95%: **{:.2}**pp, 98%: **{:.2}**pp, 99%: **{:.2}**pp, 100%: **{:.2}**pp",
@@ -98,15 +111,73 @@ pub fn beatmap_embed<'a>(
             ),
             false,
         )
+        .field("Information", diff.format_info(m, mods, None), false)
+        // .description(beatmap_description(b))
     }))
-    .field("Information", diff.format_info(m, mods, b), false)
-    .description(beatmap_description(b))
-    .footer(|f| {
-        if info.is_none() && mods != Mods::NOMOD {
-            f.text("Star difficulty not reflecting mods applied.");
-        }
-        f
-    })
+}
+
+// Some helper functions here
+
+/// Create a properly formatted beatmap title, in the `Artist - Title [Difficulty] +mods` format.
+fn beatmap_title(
+    artist: impl AsRef<str>,
+    title: impl AsRef<str>,
+    difficulty: impl AsRef<str>,
+    mods: Mods,
+) -> String {
+    let mod_str = if mods == Mods::NOMOD {
+        "".to_owned()
+    } else {
+        format!(" {}", mods)
+    };
+    MessageBuilder::new()
+        .push_bold_safe(artist.as_ref())
+        .push(" - ")
+        .push_bold_safe(title.as_ref())
+        .push(" [")
+        .push_bold_safe(difficulty.as_ref())
+        .push("]")
+        .push(&mod_str)
+        .build()
+}
+
+pub fn beatmap_embed<'a>(
+    b: &'_ Beatmap,
+    m: Mode,
+    mods: Mods,
+    info: Option<BeatmapInfoWithPP>,
+    c: &'a mut CreateEmbed,
+) -> &'a mut CreateEmbed {
+    let diff = b
+        .difficulty
+        .apply_mods(mods, info.map(|(v, _)| v.stars as f64));
+    c.title(beatmap_title(&b.artist, &b.title, &b.difficulty_name, mods))
+        .author(|a| {
+            a.name(&b.creator)
+                .url(format!("https://osu.ppy.sh/users/{}", b.creator_id))
+                .icon_url(format!("https://a.ppy.sh/{}", b.creator_id))
+        })
+        .url(b.link())
+        .image(b.cover_url())
+        .color(0xffb6c1)
+        .fields(info.map(|(_, pp)| {
+            (
+                "Calculated pp",
+                format!(
+                    "95%: **{:.2}**pp, 98%: **{:.2}**pp, 99%: **{:.2}**pp, 100%: **{:.2}**pp",
+                    pp[0], pp[1], pp[2], pp[3]
+                ),
+                false,
+            )
+        }))
+        .field("Information", diff.format_info(m, mods, b), false)
+        .description(beatmap_description(b))
+        .footer(|f| {
+            if info.is_none() && mods != Mods::NOMOD {
+                f.text("Star difficulty not reflecting mods applied.");
+            }
+            f
+        })
 }
 
 const MAX_DIFFS: usize = 25 - 4;

youmubot-osu/src/discord/hook.rs

@@ -5,7 +5,7 @@ use crate::{
 };
 use lazy_static::lazy_static;
 use regex::Regex;
-use serenity::{model::channel::Message, utils::MessageBuilder};
+use serenity::{builder::CreateEmbed, model::channel::Message, utils::MessageBuilder};
 use std::str::FromStr;
 use youmubot_prelude::*;
 
@@ -23,6 +23,101 @@ lazy_static! {
     ).unwrap();
 }
 
+pub fn dot_osu_hook<'a>(
+    ctx: &'a Context,
+    msg: &'a Message,
+) -> std::pin::Pin<Box<dyn future::Future<Output = Result<()>> + Send + 'a>> {
+    Box::pin(async move {
+        if msg.author.bot {
+            return Ok(());
+        }
+
+        // Take all the .osu attachments
+        let mut osu_embeds = msg
+            .attachments
+            .iter()
+            .filter(
+                |a| a.filename.ends_with(".osu") && a.size < 1024 * 1024, /* 1mb */
+            )
+            .map(|attachment| {
+                let url = attachment.url.clone();
+
+                async move {
+                    let data = ctx.data.read().await;
+                    let oppai = data.get::<BeatmapCache>().unwrap();
+                    let (beatmap, _) = oppai.download_beatmap_from_url(&url, None).await.ok()?;
+                    let embed_fn = crate::discord::embeds::beatmap_offline_embed(
+                        &beatmap,
+                        Mode::from(beatmap.content.mode as u8), /*For now*/
+                        msg.content.trim().parse().unwrap_or(Mods::NOMOD),
+                    )
+                    .ok()?;
+
+                    let mut create_embed = CreateEmbed::default();
+                    embed_fn(&mut create_embed);
+
+                    Some(create_embed)
+                }
+            })
+            .collect::<stream::FuturesUnordered<_>>()
+            .filter_map(|v| future::ready(v))
+            .collect::<Vec<_>>()
+            .await;
+
+        let osz_embeds = msg
+            .attachments
+            .iter()
+            .filter(
+                |a| a.filename.ends_with(".osz") && a.size < 20 * 1024 * 1024, /* 20mb */
+            )
+            .map(|attachment| {
+                let url = attachment.url.clone();
+                async move {
+                    let data = ctx.data.read().await;
+                    let oppai = data.get::<BeatmapCache>().unwrap();
+                    let beatmaps = oppai.download_osz_from_url(&url).await.pls_ok()?;
+                    Some(
+                        beatmaps
+                            .into_iter()
+                            .filter_map(|beatmap| {
+                                let embed_fn = crate::discord::embeds::beatmap_offline_embed(
+                                    &beatmap,
+                                    Mode::from(beatmap.content.mode as u8), /*For now*/
+                                    msg.content.trim().parse().unwrap_or(Mods::NOMOD),
+                                )
+                                .pls_ok()?;
+
+                                let mut create_embed = CreateEmbed::default();
+                                embed_fn(&mut create_embed);
+                                Some(create_embed)
+                            })
+                            .collect::<Vec<_>>(),
+                    )
+                }
+            })
+            .collect::<stream::FuturesUnordered<_>>()
+            .filter_map(|v| future::ready(v))
+            .filter(|v| future::ready(v.len() > 0))
+            .collect::<Vec<_>>()
+            .await
+            .concat();
+        osu_embeds.extend(osz_embeds);
+
+        if osu_embeds.len() > 0 {
+            msg.channel_id
+                .send_message(ctx, |f| {
+                    f.reference_message(msg)
+                        .content(format!("{} attached beatmaps found", osu_embeds.len()))
+                        .add_embeds(osu_embeds)
+                })
+                .await
+                .ok();
+        }
+
+        Ok(())
+    })
+}
+
 pub fn hook<'a>(
     ctx: &'a Context,
     msg: &'a Message,

youmubot-osu/src/discord/mod.rs

@@ -31,7 +31,7 @@ mod server_rank;
 use db::OsuUser;
 use db::{OsuLastBeatmap, OsuSavedUsers, OsuUserBests};
 use embeds::{beatmap_embed, score_embed, user_embed};
-pub use hook::hook;
+pub use hook::{dot_osu_hook, hook};
 use server_rank::{SERVER_RANK_COMMAND, UPDATE_LEADERBOARD_COMMAND};
 
 /// The osu! client.

youmubot-osu/src/discord/oppai_cache.rs

@@ -1,5 +1,7 @@
 use crate::mods::Mods;
+use osuparse::MetadataSection;
 use rosu_pp::{Beatmap, BeatmapExt};
+use std::io::Read;
 use std::sync::Arc;
 use youmubot_db_sql::{models::osu as models, Pool};
 use youmubot_prelude::*;
@@ -7,8 +9,9 @@ use youmubot_prelude::*;
 /// the information collected from a download/Oppai request.
 #[derive(Debug)]
 pub struct BeatmapContent {
-    id: u64,
-    content: Arc<Beatmap>,
+    id: Option<u64>,
+    pub metadata: MetadataSection,
+    pub content: Arc<Beatmap>,
 }
 
 /// the output of "one" oppai run.
@@ -128,25 +131,82 @@ impl BeatmapCache {
         BeatmapCache { client, pool }
     }
 
-    async fn download_beatmap(&self, id: u64) -> Result<BeatmapContent> {
-        let content = self
+    fn parse_beatmap(content: impl AsRef<str>, id: Option<u64>) -> Result<BeatmapContent> {
+        let content = content.as_ref();
+        let metadata = osuparse::parse_beatmap(content)
+            .map_err(|e| Error::msg(format!("Cannot parse metadata: {:?}", e)))?
+            .metadata;
+        Ok(BeatmapContent {
+            id,
+            metadata,
+            content: Arc::new(Beatmap::parse(content.as_bytes())?),
+        })
+    }
+
+    /// Downloads the given osz and try to parse every osu file in there (limited to <1mb files)
+    pub async fn download_osz_from_url(
+        &self,
+        url: impl reqwest::IntoUrl,
+    ) -> Result<Vec<BeatmapContent>> {
+        let osz = self
             .client
             .borrow()
             .await?
-            .get(&format!("https://osu.ppy.sh/osu/{}", id))
+            .get(url)
             .send()
             .await?
             .bytes()
             .await?;
-        let bm = BeatmapContent {
-            id,
-            content: Arc::new(Beatmap::parse(content.as_ref())?),
-        };
+
+        let mut osz = zip::read::ZipArchive::new(std::io::Cursor::new(osz.as_ref()))?;
+        let osu_files = osz.file_names().map(|v| v.to_owned()).collect::<Vec<_>>();
+        let osu_files = osu_files
+            .into_iter()
+            .filter(|n| n.ends_with(".osu"))
+            .filter_map(|v| {
+                let mut v = osz.by_name(&v[..]).ok()?;
+                if v.size() > 1024 * 1024
+                /*1mb*/
+                {
+                    return None;
+                };
+                let mut content = String::new();
+                v.read_to_string(&mut content).pls_ok()?;
+                Self::parse_beatmap(content, None).pls_ok()
+            })
+            .collect::<Vec<_>>();
+        Ok(osu_files)
+    }
+
+    /// Downloads the beatmap from an URL and returns it.
+    /// Does not deal with any caching.
+    pub async fn download_beatmap_from_url(
+        &self,
+        url: impl reqwest::IntoUrl,
+        id: Option<u64>,
+    ) -> Result<(BeatmapContent, String)> {
+        let content = self
+            .client
+            .borrow()
+            .await?
+            .get(url)
+            .send()
+            .await?
+            .text()
+            .await?;
+        let bm = Self::parse_beatmap(&content, id)?;
+        Ok((bm, content))
+    }
+
+    async fn download_beatmap(&self, id: u64) -> Result<BeatmapContent> {
+        let (bm, content) = self
+            .download_beatmap_from_url(&format!("https://osu.ppy.sh/osu/{}", id), Some(id))
+            .await?;
 
         let mut bc = models::CachedBeatmapContent {
             beatmap_id: id as i64,
             cached_at: chrono::Utc::now(),
-            content: content.as_ref().to_owned(),
+            content: content.into_bytes(),
         };
         bc.store(&self.pool).await?;
         Ok(bm)
@@ -155,12 +215,7 @@ impl BeatmapCache {
     async fn get_beatmap_db(&self, id: u64) -> Result<Option<BeatmapContent>> {
         Ok(models::CachedBeatmapContent::by_id(id as i64, &self.pool)
             .await?
-            .map(|v| {
-                Ok(BeatmapContent {
-                    id,
-                    content: Arc::new(Beatmap::parse(&v.content[..])?),
-                }) as Result<_>
-            })
+            .map(|v| Self::parse_beatmap(String::from_utf8(v.content)?, Some(id)))
             .transpose()?)
     }
 

youmubot-osu/src/models/mod.rs

@@ -130,34 +130,55 @@ impl Difficulty {
     }
 
     /// Format the difficulty info into a short summary.
-    pub fn format_info(&self, mode: Mode, mods: Mods, original_beatmap: &Beatmap) -> String {
-        let is_not_ranked = !matches!(original_beatmap.approval, ApprovalStatus::Ranked(_));
-        let three_lines = is_not_ranked;
+    pub fn format_info<'a>(
+        &self,
+        mode: Mode,
+        mods: Mods,
+        original_beatmap: impl Into<Option<&'a Beatmap>> + 'a,
+    ) -> String {
+        let original_beatmap = original_beatmap.into();
+        let is_not_ranked = !matches!(
+            original_beatmap.map(|v| v.approval),
+            Some(ApprovalStatus::Ranked(_))
+        );
+        let three_lines = original_beatmap.is_some() && is_not_ranked;
         let bpm = (self.bpm * 100.0).round() / 100.0;
         MessageBuilder::new()
-            .push(format!(
-                "[[Link]]({}) [[DL]]({}) [[Alt]]({}) (`{}`)",
-                original_beatmap.link(),
-                original_beatmap.download_link(false),
-                original_beatmap.download_link(true),
-                original_beatmap.short_link(Some(mode), Some(mods))
-            ))
+            .push(
+                original_beatmap
+                    .map(|original_beatmap| {
+                        format!(
+                            "[[Link]]({}) [[DL]]({}) [[Alt]]({}) (`{}`)",
+                            original_beatmap.link(),
+                            original_beatmap.download_link(false),
+                            original_beatmap.download_link(true),
+                            original_beatmap.short_link(Some(mode), Some(mods))
+                        )
+                    })
+                    .unwrap_or("**Uploaded**".to_owned()),
+            )
             .push(if three_lines { "\n" } else { ", " })
             .push_bold(format!("{:.2}⭐", self.stars))
             .push(", ")
             .push(
-                original_beatmap
-                    .difficulty
-                    .max_combo
+                self.max_combo
                     .map(|c| format!("max **{}x**, ", c))
                     .unwrap_or_else(|| "".to_owned()),
             )
             .push(if is_not_ranked {
-                format!("status **{}**, mode ", original_beatmap.approval)
+                format!(
+                    "status **{}**, mode ",
+                    original_beatmap
+                        .map(|v| v.approval)
+                        .unwrap_or(ApprovalStatus::WIP)
+                )
             } else {
                 "".to_owned()
             })
-            .push_bold_line(format_mode(mode, original_beatmap.mode))
+            .push_bold_line(format_mode(
+                mode,
+                original_beatmap.map(|v| v.mode).unwrap_or(mode),
+            ))
             .push("CS")
             .push_bold(format!("{:.1}", self.cs))
             .push(", AR")