1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
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();

        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().cloned());

        // 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 response.result.exception_details.is_some() {
            warn!(
                "[Runtime.evaluate] Caught exception while evaluating script: {:#?}",
                response
            );
            return Ok(());
        }

        Ok(())
    }
}