feat: apply userscripts to downloaded cases

Cargo.toml

@@ -2,7 +2,7 @@
 name = "aaoffline"
 description = "Downloads cases from Ace Attorney Online to be playable offline"
 repository = "https://github.com/falko17/aaoffline"
-version = "1.2.0"
+version = "1.3.0"
 edition = "2021"
 license = "MIT"
 authors = ["Falko Galperin <[email protected]>"]
@@ -50,6 +50,7 @@ too_many_arguments = "allow"
 assert_cmd = "2.0.16"
 glob = "0.3.1"
 headless_chrome = "1.0.15"
+maplit = "1.0.2"
 rstest = "0.24.0"
 rstest_reuse = "0.7.0"
 tempfile = "3.14.0"

README.md

@@ -8,6 +8,7 @@ Downloads cases from [Ace Attorney Online](https://aaonline.fr) to be playable o
 - Use parallel downloads to download case data quickly.
 - Download multiple cases at once.
 - Use the `-1` flag to compile the case into a single HTML file, without the need for a separate assets folder.
+- Apply [userscripts](https://aaonline.fr/forum/viewtopic.php?t=13534) to get a better layout, keyboard controls, and a backlog.
 - Choose a specific version of the Ace Attorney Online player (e.g., if a case only works with an older version).
 - Automatically remove photobucket watermarks from downloaded assets.
 
@@ -33,7 +34,8 @@ By default, the case will be put into a directory with the case title as its nam
 The downloaded case can then be played by opening the `index.html` file in the output directory—all case assets are put in the `assets` directory, so if you want to move this downloaded case somewhere else, you'll need to move the `assets` along with it.
 Alternatively, you can pass the `-1` flag to aaoffline, which causes the case to be compiled into a single (large) HTML file, with the assets encoded as data URLs instead of being put into separate files. (Warning: Browsers may not like HTML files very much that are multiple dozens of megabytes large. Your mileage may vary.)
 
-There are some additional parameters you can set, such as `--concurrent-downloads` to choose a different number of parallel downloads to use[^2], or `--player-version` to choose a specific commit of the player.
+There are some additional parameters you can set, such as `--concurrent-downloads` to choose a different number of parallel downloads to use[^2], `--player-version` to choose a specific commit of the player, or `--with-userscripts` to apply [userscripts](https://aaonline.fr/forum/viewtopic.php?t=13534).
+
 To get an overview of available options, just run `aaoffline --help`.
 
 ## Building / Installing

src/args.rs

@@ -2,11 +2,12 @@
 
 use anyhow::Result;
 
-use clap::{command, Parser, ValueEnum};
+use clap::{command, error::ErrorKind, CommandFactory, Parser, ValueEnum};
 #[cfg(debug_assertions)]
 use clap_verbosity_flag::DebugLevel;
 #[cfg(not(debug_assertions))]
 use clap_verbosity_flag::InfoLevel;
+use itertools::Itertools;
 use serde::Serialize;
 
 use std::path::PathBuf;
@@ -61,6 +62,21 @@ pub(crate) struct Args {
     #[arg(short('1'), long, default_value_t = false)]
     pub(crate) one_html_file: bool,
 
+    /// Whether to apply any userscripts to the downloaded case. Can be passed multiple times.
+    ///
+    /// Scripts were created by Time Axis, with only the expanded keyboard controls written by me,
+    /// building on Time Axis' basic keyboard controls script.
+    /// (These options may change in the future when some scripts are consolidated).
+    #[arg(
+        short('u'),
+        long,
+        num_args(0..=1),
+        default_missing_value("all"),
+        require_equals(true),
+        value_enum,
+    )]
+    pub(crate) with_userscripts: Vec<Userscripts>,
+
     /// How many concurrent downloads to use.
     #[arg(short('j'), long, default_value_t = 5)]
     pub(crate) concurrent_downloads: usize,
@@ -139,6 +155,65 @@ pub(crate) enum DownloadSequence {
     Ask,
 }
 
+/// Whether to apply any userscripts to the downloaded case.
+#[derive(Debug, ValueEnum, Clone, Serialize, Default, PartialEq, Eq)]
+#[serde(rename_all = "kebab-case")]
+pub(crate) enum Userscripts {
+    /// Apply all userscripts.
+    All,
+
+    /// Changes the fonts of nametags to use a proper pixelized font.
+    AltNametag,
+    /// Adds a backlog button to see past dialog.
+    Backlog,
+    /// Improves the layout (e.g., enlarging and centering the main screens).
+    BetterLayout,
+    /// Adds extensive keyboard controls. See the top of the file at
+    /// <https://gist.github.com/falko17/965207b1f1f0496ff5f0cb41d8e827f2#file-aaokeyboard-user-js>
+    /// to get an overview of available controls.
+    KeyboardControls,
+
+    /// Apply no userscript.
+    #[default]
+    None,
+}
+
+impl Userscripts {
+    /// Returns the URLs pointing to the corresponding userscripts.
+    pub(crate) fn urls(&self) -> Vec<&str> {
+        match self {
+            Self::AltNametag => vec!["https://beyondtimeaxis.github.io/misc/aaoaltnametags.user.js"],
+            Self::Backlog => vec!["https://beyondtimeaxis.github.io/misc/aaobacklog.user.js"],
+            Self::BetterLayout => vec!["https://beyondtimeaxis.github.io/misc/aaobetterlayout.user.js"],
+            Self::KeyboardControls => vec!["https://gist.github.com/falko17/965207b1f1f0496ff5f0cb41d8e827f2/raw/aaokeyboard.user.js"],
+            Self::All => [Self::AltNametag, Self::Backlog, Self::BetterLayout, Self::KeyboardControls].iter().flat_map(Self::urls).collect(),
+            Self::None => vec![]
+        }
+    }
+
+    /// Returns all URLs belonging to the given collection of [scripts].
+    pub(crate) fn all_urls(scripts: &[Self]) -> Vec<&str> {
+        scripts.iter().flat_map(|x| x.urls()).unique().collect()
+    }
+
+    /// Ensures that the given [scripts] are a valid combination.
+    pub(crate) fn validate_combination(scripts: &[Self]) -> Result<(), clap::Error> {
+        if (scripts.contains(&Self::All)) && scripts.len() > 1 {
+            Err(Args::command().error(
+                ErrorKind::ArgumentConflict,
+                "Can't specify any other scripts when including all of them anyway",
+            ))
+        } else if scripts.contains(&Self::None) && scripts.len() > 1 {
+            Err(Args::command().error(
+                ErrorKind::ArgumentConflict,
+                "Can't specify any other scripts when including none",
+            ))
+        } else {
+            Ok(())
+        }
+    }
+}
+
 impl Args {
     /// Parses the given [case] into its ID.
     fn accept_case(case: &str) -> Result<u32, String> {

src/data/player.rs

@@ -1,5 +1,6 @@
 //! Contains data model related to the case player and its scripts.
 
+use crate::args::Userscripts;
 use crate::constants::{re, AAONLINE_BASE, BITBUCKET_URL};
 use crate::download::Download;
 use crate::transform::php;
@@ -403,7 +404,42 @@ impl Player {
     pub(crate) async fn retrieve_scripts(&mut self, pb: &ProgressBar) -> Result<()> {
         self.scripts
             .retrieve_player_scripts(&self.site_data, pb, Self::transform_module)
-            .await?;
+            .await
+    }
+
+    /// Retrieves the userscripts and appends them to the player scripts.
+    pub(crate) async fn retrieve_userscripts(&mut self, pb: &ProgressBar) -> Result<()> {
+        const HTML_END: &str = "</html>";
+        let urls = Userscripts::all_urls(&self.scripts.ctx.args.with_userscripts);
+        let client = &self.scripts.ctx.client;
+        let userscripts = stream::iter(urls)
+            .map(|url| async move {
+                debug!("Downloading userscript {url}...");
+                pb.inc(1);
+                client
+                    .get(url)
+                    .send()
+                    .await
+                    .context("Could not download userscript.")?
+                    .error_for_status()
+                    .context("Userscript seems to be inaccessible.")?
+                    .text()
+                    .await
+                    .context("Script could not be decoded as text")
+            })
+            .buffer_unordered(self.scripts.ctx.args.concurrent_downloads)
+            .collect::<Vec<_>>()
+            .await
+            .into_iter()
+            .flatten()
+            .join("\n\n");
+        let content = self.content.as_mut().expect("player must be present");
+        let html_end = content
+            .rfind(HTML_END)
+            .expect("end of player must be present");
+        let replacement =
+            format!("<script type=\"text/javascript\">{userscripts}</script>\n{HTML_END}");
+        content.replace_range(html_end..html_end + HTML_END.len(), &replacement);
         Ok(())
     }
 

src/main.rs

@@ -10,6 +10,7 @@ mod middleware;
 pub(crate) mod transform;
 
 use anyhow::{Context, Result};
+use args::Userscripts;
 use clap::Parser;
 use colored::Colorize;
 use data::case::{Case, Sequence};
@@ -140,7 +141,7 @@ impl MainContext {
     /// using the given [ctx] for the arguments.
     fn show_step_ctx(&self, step: u8, text: &str, ctx: &GlobalContext) {
         self.pb
-            .set_message(format!("{} {text}", format!("[{step}/7]").dimmed()));
+            .set_message(format!("{} {text}", format!("[{step}/8]").dimmed()));
         if !Self::should_hide_pb(&ctx.args) {
             self.pb.enable_steady_tick(Duration::from_millis(50));
         }
@@ -389,6 +390,23 @@ impl MainContext {
         self.clean_on_fail(result).await
     }
 
+    /// Retrieves the userscripts and appends them to the player.
+    async fn append_userscripts(&mut self) -> Result<()> {
+        let urls = Userscripts::all_urls(&self.ctx().args.with_userscripts);
+        if urls.is_empty() {
+            return Ok(());
+        }
+        let pb = self.add_progress(urls.len() as u64);
+        let result = self
+            .player
+            .as_mut()
+            .unwrap()
+            .retrieve_userscripts(&pb)
+            .await;
+        self.finish_progress(&pb, "Userscripts retrieved.");
+        self.clean_on_fail(result).await
+    }
+
     /// Transforms the player blocks for the given [case] to point to offline assets.
     async fn transform_player_blocks(&mut self, case: &Case) -> Result<()> {
         let result = self.player.as_mut().unwrap().transform_player(case);
@@ -437,6 +455,7 @@ impl MainContext {
 async fn main() -> Result<()> {
     setup_panic!();
     let args = args::Args::parse();
+    Userscripts::validate_combination(&args.with_userscripts)?;
     let original_output = args.output.clone();
     let one_file = args.one_html_file;
     env_logger::builder()
@@ -567,12 +586,15 @@ async fn main() -> Result<()> {
     ctx.show_step(6, "Retrieving additional external player sources...");
     ctx.retrieve_player_sources().await?;
 
+    ctx.show_step(7, "Applying userscripts...");
+    ctx.append_userscripts().await?;
+
     let original_state = ctx.player.as_ref().unwrap().save();
     let mut output_path: &PathBuf = &PathBuf::new();
     for case in cases {
         // Need to reset transformed player.
         ctx.show_step(
-            7,
+            8,
             &format!(
                 "Writing case \"{}\" to disk...",
                 case.case_information.title

tests/integration_test.rs

@@ -1,6 +1,7 @@
 use std::{
     error::Error,
     fmt::Display,
+    fs,
     path::PathBuf,
     sync::{Arc, Mutex, Weak},
     time::Duration,
@@ -20,6 +21,7 @@ use headless_chrome::{
     Browser, LaunchOptionsBuilder,
 };
 use itertools::Itertools;
+use maplit::hashmap;
 use rstest::{fixture, rstest};
 use rstest_reuse::{apply, template};
 use tempfile::{tempdir, TempDir};
@@ -178,11 +180,12 @@ fn verify_with_browser_common(
     tab.enable_log()?;
     let listener: Arc<JsListener> = Arc::new(JsListener::default());
     tab.add_event_listener(listener.clone())?;
-    tab.navigate_to(&format!(
+    let url = format!(
         "file://{path}{}.html{}",
         if append_index { "/index" } else { "" },
         query.unwrap_or_default()
-    ))?;
+    );
+    tab.navigate_to(&url)?;
     tab.wait_until_navigated()?;
 
     // We click the "Start" button and wait a little while.
@@ -336,6 +339,80 @@ fn test_multi(mut cmd: Cmd) {
     drop(cmd);
 }
 
+#[rstest]
+fn test_userscripts(
+    mut cmd: Cmd,
+    #[values(
+        "",
+        "=all",
+        "=none",
+        "=alt-nametag",
+        "=backlog",
+        "=better-layout",
+        "=keyboard-controls"
+    )]
+    userscripts: &str,
+    #[values(true, false)] one_file: bool,
+) {
+    use regex::Regex;
+
+    cmd.with_tmp_output(one_file);
+    if one_file {
+        cmd.cmd.arg("-1");
+    }
+    cmd.cmd.arg(PSYCHE_LOCK_TEST);
+    cmd.cmd.arg(String::from("-u") + userscripts);
+    println!(
+        "Command-line args: {}",
+        cmd.cmd.get_args().map(|x| x.to_str().unwrap()).join(" ")
+    );
+    cmd.cmd.assert().success();
+    // We want to both make sure that the userscript is actually present...
+    let verify_regexes = hashmap! {
+        "=alt-nametag" => r"// @name\s+AAO Alt Nametag Font",
+        "=backlog" => r"// @name\s+AAO Backlog Script",
+        "=better-layout" => r"// @name\s+AAO Better Layout Script",
+        "=keyboard-controls" => r"// @name\s+AAO Keyboard Controls \(Expanded\)",
+    };
+
+    let path = format!("{}/index.html", cmd.path_as_str(),);
+    let file = fs::read_to_string(path).unwrap();
+    if let Some(regextext) = verify_regexes.get(userscripts) {
+        let regex = Regex::new(regextext).unwrap();
+        assert!(regex.is_match(&file));
+    } else if userscripts == "=all" || userscripts.is_empty() {
+        for r in verify_regexes.values() {
+            let regex = Regex::new(r).unwrap();
+            assert!(regex.is_match(&file));
+        }
+    } else if userscripts == "=none" {
+        assert!(!file.contains("==UserScript=="));
+    } else {
+        panic!("Unknown userscripts setting");
+    }
+    // ...and that there are no errors in the scripts.
+    verify_with_browser(cmd.path_as_str(), None).unwrap();
+}
+
+#[rstest]
+fn test_invalid_userscript(
+    mut cmd: Cmd,
+    #[values(
+        "-u backlog -u all",
+        "-u -u",
+        "-u backlog -u none",
+        "-u none -u all",
+        "-u some"
+    )]
+    userscripts: &str,
+) {
+    cmd.cmd
+        .arg(PSYCHE_LOCK_TEST)
+        .args(userscripts.split(' '))
+        .assert()
+        .failure();
+}
+
 #[rstest]
 fn test_output_format(
     #[values(true, false)] one_file: bool,