electron_injector/
injector.rs

1use std::process::{Child, Command};
2use std::{fs, thread, time};
3
4use anyhow::{bail, Result};
5use log::{debug, info, warn};
6use portpicker;
7use serde_json::json;
8
9use crate::assets;
10use crate::config::Config;
11use crate::protocol::{DevtoolPage, EvaluateResponse};
12use crate::websocket::WebSocket;
13
14struct UserScript {
15    file_path: String,
16    content: String,
17}
18
19pub struct Injector {
20    config: Config,
21    port: u16,
22}
23
24impl Injector {
25    pub(crate) const INJECT_LOOP_SLEEP_MS: u64 = 1000;
26    pub(crate) const WAIT_DEBUGGING_PORT_TIMEOUT_MS: u64 = 30_000;
27
28    fn get_available_port(config: &Config) -> u16 {
29        if portpicker::is_free_tcp(config.port) {
30            info!("Using port: {}", config.port);
31            return config.port;
32        }
33
34        info!(
35            "Port {} is not available, finding another port",
36            config.port
37        );
38
39        let port = portpicker::pick_unused_port().expect("Port should be available");
40        info!("Found available port: {}", port);
41
42        port
43    }
44
45    pub fn new() -> Self {
46        // Parse CLI args
47        let config = Config::parse_auto();
48
49        // Get port
50        let port = Injector::get_available_port(&config);
51
52        Injector { config, port }
53    }
54
55    pub fn run(&self) -> Result<()> {
56        info!("Running injector");
57        debug!("{:#?}", self.config);
58
59        // Spawn child process
60        _ = self.spawn_process()?;
61
62        // Prepare prelude script
63        let prelude_script = self.get_prelude_script().unwrap_or(String::new());
64
65        // Prepare user scripts
66        let user_scripts = self.get_user_scripts();
67
68        // Create timeout duration
69        let timeout_duration = time::Duration::from_millis(self.config.timeout);
70
71        // Declare a vec to store found page ids
72        let mut found_page_ids: Vec<String> = Vec::new();
73
74        // Inject loop
75        let start_time = time::Instant::now();
76        loop {
77            // Refresh devtool pages
78            let devtool_pages = self
79                .get_devtool_pages()
80                .expect("Should be able to get devtool pages");
81
82            debug!("{:#?}", devtool_pages);
83
84            // Loop through pages
85            for page in devtool_pages {
86                if found_page_ids.contains(&page.id) {
87                    continue;
88                }
89
90                // Create WebSocket
91                let mut ws = WebSocket::connect(&page.web_socket_debugger_url)
92                    .expect("To connect to websocket");
93
94                // Inject prelude
95                if self.config.prelude {
96                    info!("Injecting prelude script (id: {})", page.id);
97                    self.evaluate(&mut ws, &prelude_script)
98                        .expect("Should be able to evaluate JS");
99                }
100
101                // Inject scripts
102                for user_script in user_scripts.iter() {
103                    // Inject using evaluate
104                    info!("Injecting script: {}", user_script.file_path);
105                    self.evaluate(&mut ws, &user_script.content)
106                        .expect("Should be able to evaluate JS");
107                }
108
109                // Save page id
110                found_page_ids.push(page.id.clone());
111            }
112
113            // Check devtool pages again
114            let updated_devtool_pages = self
115                .get_devtool_pages()
116                .expect("Should be able to get devtool pages");
117
118            // Stop if already found all pages
119            if found_page_ids.len() == updated_devtool_pages.len() {
120                info!("Stopping injection loop");
121                break;
122            }
123
124            // Timed out
125            if start_time.elapsed() >= timeout_duration {
126                bail!("Injection loop timed out");
127            }
128
129            // Sleep before next loop iteration
130            thread::sleep(time::Duration::from_millis(Self::INJECT_LOOP_SLEEP_MS));
131        }
132
133        info!("Injection success");
134        Ok(())
135    }
136
137    fn get_devtool_pages(&self) -> Result<Vec<DevtoolPage>, reqwest::Error> {
138        let url = format!("http://{}:{}/json/list", &self.config.host, &self.port);
139
140        let client = reqwest::blocking::Client::new();
141        let response = client.get(url).send()?.error_for_status()?;
142
143        let pages_response = response.json::<Vec<DevtoolPage>>()?;
144        Ok(pages_response)
145    }
146
147    fn get_prelude_script(&self) -> Option<String> {
148        // No need to load if not enabled anyways
149        if !self.config.prelude {
150            return None;
151        }
152
153        // Load from embedded file
154        let file = assets::JS::get("prelude.js").unwrap();
155        let script =
156            std::str::from_utf8(file.data.as_ref()).expect("Script should be a valid UTF-8 file");
157
158        Some(String::from(script))
159    }
160
161    fn get_user_scripts(&self) -> Vec<UserScript> {
162        let scripts: Vec<UserScript> = self
163            .config
164            .script
165            .iter()
166            .map(|s| {
167                let content =
168                    fs::read_to_string(s).expect("Should have been able to read the file");
169
170                UserScript {
171                    file_path: s.to_string(),
172                    content,
173                }
174            })
175            .collect();
176
177        scripts
178    }
179
180    fn spawn_process(&self) -> Result<Child> {
181        // Prepare args
182        let mut args = vec![format!("--remote-debugging-port={}", &self.port)];
183        args.extend(self.config.arg.iter().cloned());
184
185        // Spawn child process
186        debug!(
187            "Spawning electron app: {} (args: {:#?})",
188            &self.config.app, args
189        );
190        let child = Command::new(&self.config.app).args(args).spawn()?;
191
192        // Wait for process
193        info!("Waiting for {}ms", self.config.delay);
194        thread::sleep(time::Duration::from_millis(self.config.delay));
195
196        // Create timeout duration
197        let timeout_duration = time::Duration::from_millis(Self::WAIT_DEBUGGING_PORT_TIMEOUT_MS);
198
199        // Wait until remote debugging port is available
200        info!("Waiting for remote debugging port");
201        let start_time = time::Instant::now();
202        loop {
203            // Connected
204            if self.get_devtool_pages().is_ok() {
205                info!("Connected to remote debugging port");
206                break;
207            }
208
209            // Timed out
210            if start_time.elapsed() >= timeout_duration {
211                bail!("Unable to connect to remote debugging port");
212            }
213        }
214
215        Ok(child)
216    }
217
218    fn evaluate(&self, ws: &mut WebSocket, expression: &str) -> Result<()> {
219        // Create payload
220        // https://chromedevtools.github.io/devtools-protocol/tot/Runtime/#method-evaluate
221        let payload = json!({
222            "id": 1,
223            "method": "Runtime.evaluate",
224            "params": {
225                "expression": expression,
226                "objectGroup": "inject",
227                "includeCommandLineAPI": true,
228                "silent": true,
229                "userGesture": true,
230                "awaitPromise": true,
231            },
232        });
233
234        // Serialize payload to JSON
235        let payload_json = serde_json::to_string(&payload)?;
236
237        // Send message and get the result
238        let result_msg = ws.send_and_receive(&payload_json)?;
239        debug!("[Runtime.evaluate] Raw message: {:#?}", result_msg);
240
241        // Ignore if not a text
242        if !result_msg.is_text() {
243            warn!(
244                "[Runtime.evaluate] Unexpected result from WebSocket: {:#?}",
245                result_msg
246            );
247            return Ok(());
248        }
249
250        // Convert message to text
251        let result_json = result_msg.to_text()?;
252
253        // Parse response
254        let response: EvaluateResponse = serde_json::from_str(result_json)?;
255
256        debug!("[Runtime.evaluate] Parsed response: {:#?}", response);
257
258        // Handle exception
259        if response.result.exception_details.is_some() {
260            warn!(
261                "[Runtime.evaluate] Caught exception while evaluating script: {:#?}",
262                response
263            );
264            return Ok(());
265        }
266
267        Ok(())
268    }
269}
electron_injector/injector.rs

electron_injector/
injector.rs
