initial commit

src/assets.rs

@@ -0,0 +1,5 @@
+use rust_embed::RustEmbed;
+
+#[derive(RustEmbed)]
+#[folder = "assets/js/"]
+pub struct JS;

src/config.rs

@@ -0,0 +1,62 @@
+use std::ops::RangeInclusive;
+
+use clap::Parser;
+
+const PORT_RANGE: RangeInclusive<usize> = 1..=65535;
+
+fn validate_port(s: &str) -> Result<u16, String> {
+    let port: usize = s
+        .parse()
+        .map_err(|_| format!("`{s}` isn't a port number"))?;
+    if PORT_RANGE.contains(&port) {
+        Ok(port as u16)
+    } else {
+        Err(format!(
+            "port not in range {}-{}",
+            PORT_RANGE.start(),
+            PORT_RANGE.end()
+        ))
+    }
+}
+
+#[derive(Debug, Parser)]
+#[command(author, version, about, long_about = None)]
+pub struct Config {
+    /// Path to the electron app
+    #[arg()]
+    pub app: String,
+
+    /// Additional arg for the electron app
+    #[arg(short, long)]
+    pub arg: Vec<String>,
+
+    /// Path to the javascript file to be injected
+    #[arg(short, long)]
+    pub script: Vec<String>,
+
+    /// The remote debugging host
+    #[arg(long, default_value_t = String::from("127.0.0.1"))]
+    pub host: String,
+
+    /// The remote debugging port
+    #[arg(short, long, default_value_t = 8315, value_parser = validate_port)]
+    pub port: u16,
+
+    /// Timeout in ms for injecting scripts
+    #[arg(short, long, default_value_t = 10_000)]
+    pub timeout: u64,
+
+    /// Delay in ms to wait after spawning the process
+    #[arg(short, long, default_value_t = 10_000)]
+    pub delay: u64,
+
+    /// Enable prelude script
+    #[arg(long)]
+    pub prelude: bool,
+}
+
+impl Config {
+    pub fn parse_auto() -> Config {
+        Config::parse()
+    }
+}

src/injector.rs

@@ -0,0 +1,271 @@
+use std::process::{Child, Command};
+use std::{fs, thread, time};
+
+use anyhow::{bail, Result};
+use log::{debug, info, warn};
+use portpicker;
+use serde_json::json;
+
+use crate::assets;
+use crate::config::Config;
+use crate::protocol::{DevtoolPage, EvaluateResponse};
+use crate::websocket::WebSocket;
+
+struct UserScript {
+    file_path: String,
+    content: String,
+}
+
+pub struct Injector {
+    config: Config,
+    port: u16,
+}
+
+impl Injector {
+    pub(crate) const INJECT_LOOP_SLEEP_MS: u64 = 1000;
+    pub(crate) const WAIT_DEBUGGING_PORT_TIMEOUT_MS: u64 = 30_000;
+
+    fn get_available_port(config: &Config) -> u16 {
+        if portpicker::is_free_tcp(config.port) {
+            info!("Using port: {}", config.port);
+            return config.port;
+        }
+
+        info!(
+            "Port {} is not available, finding another port",
+            config.port
+        );
+
+        let port = portpicker::pick_unused_port().expect("Port should be available");
+        info!("Found available port: {}", port);
+
+        port
+    }
+
+    pub fn new() -> Self {
+        // Parse CLI args
+        let config = Config::parse_auto();
+
+        // Get port
+        let port = Injector::get_available_port(&config);
+
+        Injector { config, port }
+    }
+
+    pub fn run(&self) -> Result<()> {
+        info!("Running injector");
+        debug!("{:#?}", self.config);
+
+        // Spawn child process
+        _ = self.spawn_process()?;
+
+        // Prepare prelude script
+        let prelude_script = self.get_prelude_script().unwrap_or(String::new());
+
+        // Prepare user scripts
+        let user_scripts = self.get_user_scripts();
+
+        // Create timeout duration
+        let timeout_duration = time::Duration::from_millis(self.config.timeout);
+
+        // Declare a vec to store found page ids
+        let mut found_page_ids: Vec<String> = Vec::new();
+
+        // Inject loop
+        let start_time = time::Instant::now();
+        loop {
+            // Refresh devtool pages
+            let devtool_pages = self
+                .get_devtool_pages()
+                .expect("Should be able to get devtool pages");
+
+            debug!("{:#?}", devtool_pages);
+
+            // Loop through pages
+            for page in devtool_pages {
+                if found_page_ids.contains(&page.id) {
+                    continue;
+                }
+
+                // Create WebSocket
+                let mut ws = WebSocket::connect(&page.web_socket_debugger_url)
+                    .expect("To connect to websocket");
+
+                // Inject prelude
+                if self.config.prelude {
+                    info!("Injecting prelude script (id: {})", page.id);
+                    _ = self
+                        .evaluate(&mut ws, &prelude_script)
+                        .expect("Should be able to evaluate JS");
+                }
+
+                // Inject scripts
+                for user_script in user_scripts.iter() {
+                    // Inject using evaluate
+                    info!("Injecting script: {}", user_script.file_path);
+                    _ = self
+                        .evaluate(&mut ws, &user_script.content)
+                        .expect("Should be able to evaluate JS");
+                }
+
+                // Save page id
+                found_page_ids.push(page.id.clone());
+            }
+
+            // Check devtool pages again
+            let updated_devtool_pages = self
+                .get_devtool_pages()
+                .expect("Should be able to get devtool pages");
+
+            // Stop if already found all pages
+            if found_page_ids.len() == updated_devtool_pages.len() {
+                info!("Stopping injection loop");
+                break;
+            }
+
+            // Timed out
+            if start_time.elapsed() >= timeout_duration {
+                bail!("Injection loop timed out");
+            }
+
+            // Sleep before next loop iteration
+            thread::sleep(time::Duration::from_millis(Self::INJECT_LOOP_SLEEP_MS));
+        }
+
+        info!("Injection success");
+        Ok(())
+    }
+
+    fn get_devtool_pages(&self) -> Result<Vec<DevtoolPage>, reqwest::Error> {
+        let url = format!("http://{}:{}/json/list", &self.config.host, &self.port);
+
+        let client = reqwest::blocking::Client::new();
+        let response = client.get(url).send()?.error_for_status()?;
+
+        let pages_response = response.json::<Vec<DevtoolPage>>()?;
+        Ok(pages_response)
+    }
+
+    fn get_prelude_script(&self) -> Option<String> {
+        // No need to load if not enabled anyways
+        if !self.config.prelude {
+            return None;
+        }
+
+        // Load from embedded file
+        let file = assets::JS::get("prelude.js").unwrap();
+        let script =
+            std::str::from_utf8(file.data.as_ref()).expect("Script should be a valid UTF-8 file");
+
+        Some(String::from(script))
+    }
+
+    fn get_user_scripts(&self) -> Vec<UserScript> {
+        let scripts: Vec<UserScript> = self
+            .config
+            .script
+            .iter()
+            .map(|s| {
+                let content =
+                    fs::read_to_string(s).expect("Should have been able to read the file");
+
+                UserScript {
+                    file_path: s.to_string(),
+                    content,
+                }
+            })
+            .collect();
+
+        return scripts;
+    }
+
+    fn spawn_process(&self) -> Result<Child> {
+        // Prepare args
+        let mut args = vec![format!("--remote-debugging-port={}", &self.port)];
+        args.extend(self.config.arg.iter().map(|a| a.clone()));
+
+        // Spawn child process
+        debug!(
+            "Spawning electron app: {} (args: {:#?})",
+            &self.config.app, args
+        );
+        let child = Command::new(&self.config.app).args(args).spawn()?;
+
+        // Wait for process
+        info!("Waiting for {}ms", self.config.delay);
+        thread::sleep(time::Duration::from_millis(self.config.delay));
+
+        // Create timeout duration
+        let timeout_duration = time::Duration::from_millis(Self::WAIT_DEBUGGING_PORT_TIMEOUT_MS);
+
+        // Wait until remote debugging port is available
+        info!("Waiting for remote debugging port");
+        let start_time = time::Instant::now();
+        loop {
+            // Connected
+            if self.get_devtool_pages().is_ok() {
+                info!("Connected to remote debugging port");
+                break;
+            }
+
+            // Timed out
+            if start_time.elapsed() >= timeout_duration {
+                bail!("Unable to connect to remote debugging port");
+            }
+        }
+
+        Ok(child)
+    }
+
+    fn evaluate(&self, ws: &mut WebSocket, expression: &str) -> Result<()> {
+        // Create payload
+        // https://chromedevtools.github.io/devtools-protocol/tot/Runtime/#method-evaluate
+        let payload = json!({
+            "id": 1,
+            "method": "Runtime.evaluate",
+            "params": {
+                "expression": expression,
+                "objectGroup": "inject",
+                "includeCommandLineAPI": true,
+                "silent": true,
+                "userGesture": true,
+                "awaitPromise": true,
+            },
+        });
+
+        // Serialize payload to JSON
+        let payload_json = serde_json::to_string(&payload)?;
+
+        // Send message and get the result
+        let result_msg = ws.send_and_receive(&payload_json)?;
+        debug!("[Runtime.evaluate] Raw message: {:#?}", result_msg);
+
+        // Ignore if not a text
+        if !result_msg.is_text() {
+            warn!(
+                "[Runtime.evaluate] Unexpected result from WebSocket: {:#?}",
+                result_msg
+            );
+            return Ok(());
+        }
+
+        // Convert message to text
+        let result_json = result_msg.to_text()?;
+
+        // Parse response
+        let response: EvaluateResponse = serde_json::from_str(result_json)?;
+
+        debug!("[Runtime.evaluate] Parsed response: {:#?}", response);
+
+        // Handle exception
+        if let Some(_) = response.result.exception_details {
+            warn!(
+                "[Runtime.evaluate] Caught exception while evaluating script: {:#?}",
+                response
+            );
+            return Ok(());
+        }
+
+        Ok(())
+    }
+}

src/lib.rs

@@ -0,0 +1,5 @@
+pub mod assets;
+pub mod config;
+pub mod injector;
+mod protocol;
+mod websocket;

src/main.rs

@@ -0,0 +1,10 @@
+use electron_injector::injector::Injector;
+
+fn main() {
+    // Setup logging
+    pretty_env_logger::init();
+
+    // Run the injector
+    let injector = Injector::new();
+    injector.run().unwrap();
+}

src/protocol.rs

@@ -0,0 +1,36 @@
+use serde::Deserialize;
+use serde_json::Value;
+
+#[derive(Debug, Deserialize)]
+pub struct DevtoolPage {
+    #[serde(rename = "description")]
+    pub description: String,
+    #[serde(rename = "devtoolsFrontendUrl")]
+    pub devtools_frontend_url: String,
+    #[serde(rename = "id")]
+    pub id: String,
+    #[serde(rename = "title")]
+    pub title: String,
+    #[serde(rename = "type")]
+    pub r#type: String,
+    #[serde(rename = "url")]
+    pub url: String,
+    #[serde(rename = "webSocketDebuggerUrl")]
+    pub web_socket_debugger_url: String,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct EvaluateResponse {
+    #[serde(rename = "id")]
+    pub id: i32,
+    #[serde(rename = "result")]
+    pub result: EvaluateResult,
+}
+
+#[derive(Debug, Deserialize)]
+pub struct EvaluateResult {
+    #[serde(rename = "result")]
+    pub result: Value,
+    #[serde(rename = "exceptionDetails")]
+    pub exception_details: Option<Value>,
+}

src/websocket.rs

@@ -0,0 +1,54 @@
+use std::net::TcpStream;
+
+use log::{debug, info};
+use tungstenite::{stream::MaybeTlsStream, Message};
+
+pub struct WebSocket {
+    address: String,
+    socket: tungstenite::WebSocket<MaybeTlsStream<TcpStream>>,
+}
+
+impl WebSocket {
+    pub fn connect(address: &str) -> Result<Self, tungstenite::Error> {
+        let url = url::Url::parse(&address).expect("Should be a valid address");
+
+        let (socket, response) = tungstenite::connect(url)?;
+        info!("WebSocket connected (status: {})", response.status());
+
+        debug!("Response headers: {:#?}", response.headers());
+
+        Ok(WebSocket {
+            address: String::from(address),
+            socket,
+        })
+    }
+
+    pub fn send(&mut self, msg: &str) -> Result<(), tungstenite::Error> {
+        self.socket
+            .write_message(Message::Text(String::from(msg)))?;
+
+        Ok(())
+    }
+
+    pub fn receive(&mut self) -> Result<Message, tungstenite::Error> {
+        let msg = self.socket.read_message()?;
+
+        Ok(msg)
+    }
+
+    pub fn send_and_receive(&mut self, msg: &str) -> Result<Message, tungstenite::Error> {
+        self.send(msg)?;
+        self.receive()
+    }
+
+    pub fn close(&mut self) -> Result<(), tungstenite::Error> {
+        self.socket.close(None)
+    }
+}
+
+impl Drop for WebSocket {
+    fn drop(&mut self) {
+        debug!("Closing WebSocket (address: {:?})", self.address);
+        _ = self.close();
+    }
+}